Search code examples
node.jstypescriptjestjsopenai-api

Using an asyncIterator inside a unit test never resolves with values


I have a readable stream from chatGPT that I'm trying to assert on using an asyncIterator and then reading those values and asserting on them. The jest test runs fine but there is an error in my logic that is stopping it from resolving after each iteration.

What am I missing?

Thanks Jimi

var util = require("util");

describe("stream", () => {
  it("should work", async () => {
    let i = -1;
    const completionStream: ChatCompletionStreamParams = {
      toReadableStream: () => {
        const mockReadable = {
          async *[Symbol.asyncIterator]() {
            const encoder = new util.TextEncoder("utf-8");
            yield encoder.encode(
              JSON.stringify({
                id: "1",
                object: "text",
                created: 1629159835,
                model: "gpt4",
                choices: [
                  {
                    index: 0,
                    delta: { content: "Hello" },
                    finish_reason: null,
                  },
                ],
              })
            ).buffer;
            yield encoder.encode(
              JSON.stringify({
                id: "1",
                object: "text",
                created: 1629159835,
                model: "gpt4",
                choices: [
                  {
                    index: 0,
                    delta: { content: "there" },
                    finish_reason: null,
                  },
                ],
              })
            ).buffer;
            yield null;
          },
          getReader: () => {
            return {
              read: async () => {
                const result = await mockReadable[
                  Symbol.asyncIterator
                ]().next();
                i++;
                if (i >= 2) {
                  return Promise.resolve({ done: true, value: null });
                }
                return Promise.resolve({ done: false, value: result?.value });
              },
            };
          },
        };
        return mockReadable;
      },
    };

    const readable = completionStream.toReadableStream().getReader();

    while (true) {
      const { done, value } = await readable.read();
      expect(done).toBe(false); // assert here
    }
  });
});

Solution

  • You made a mistake in getReader at the mockReadable[Symbol.asyncIterator]() part. It is called every time you call readable.read(), which creates a fresh async iterator from scratch, thus it keeps yielding "Hello" and never reaches "there".

    You need to save the iterator and reuse it.

    getReader: () => {
      const iterator = mockReadable[Symbol.asyncIterator]();
      return {
        read: async () => {
          const result = await iterator.next();
          return result;
        }
    
        // in fact you can just re-expose `iterator.next` as `read` like
        // read: iterator.next.bind(iterator),
      };
    }