如何用Jest为OpenAI Node SDK（streaming=true）编写集成测试

我在给使用streaming: true的OpenAI Node.js SDK写测试时遇到麻烦，浏览器里代码跑正常，但测试里拦不住OpenAI的请求，也没法通过模拟返回正确响应让代码继续执行。我想测试整个Express路由，不限制用现有Mock库，求替代方案。

测试代码 #

it("should return stream the ai output if this is the first time fetching the document", async () => {
      nock("https://api.openai.com")
        .post("/v1/chat/completions")
        .reply(200, {
          id: "chatcmpl-123",
          object: "chat.completion.chunk",
          created: 1694268190,
          model: "gpt-3.5-turbo-0613",
          system_fingerprint: "fp_44709d6fcb",
          choices: [
            {
              index: 0,
              delta: {
                role: "assistant",
                content: "I am actually not OpenAI.",
              },
            },
          ],
          finish_reason: "stop",
        });
      const response = await agent.get(`${baseRoute}/documents/${documentId}`);
      expect(response.status).toBe(200);
      console.log(response.body);
    });

/api/v2/documents/:id控制器代码 #

const stream = await openaiService.getCompletionStream(document);
response.setHeader("Transfer-Encoding", "chunked");
response.setHeader("X-Content-Type-Options", "nosniff");
for await (const chunk of stream) {
  const data = chunk.choices[0]?.delta?.content;
  if (data !== undefined) {
    response.write(data);
  }
}

response.end();

被调用的getCompletionStream函数代码 #

getCompletionStream = async (document: DocumentDtoV2) => {
    const openai = this.getOpenAiAgent();
    const completion = await openai.chat.completions.create({
      model: document.aiEngine.defaultProperties.model,
      messages: document.messages,
      max_tokens: document.aiEngine.defaultProperties.max_tokens,
      stream: true,
      temperature: document.prompt.model.customProperties.temperature
        ? document.prompt.model.customProperties.temperature
        : document.aiEngine.defaultProperties.temperature,
      top_p: document.prompt.model.customProperties.top_p
        ? document.prompt.model.customProperties.top_p
        : document.aiEngine.defaultProperties.top_p,
      frequency_penalty: (document.aiEngine.defaultProperties as ChatGPTProperties)
        .frequency_penalty,
      presence_penalty: (document.aiEngine.defaultProperties as ChatGPTProperties).presence_penalty,
    });

    return completion;
  };

核心问题分析 #

nock默认返回完整JSON响应，但OpenAI流式接口返回的是分块的Server-Sent Events(SSE)，格式为data: {JSON chunk}\n\n，直接返回单个JSON对象模拟不了流式响应，导致代码里的for await...of循环无法正确处理。

方案1：修复nock的流式响应模拟 #

修改nock的reply，模拟SSE格式的分块响应：

it("should return stream the ai output if this is the first time fetching the document", async () => {
  nock("https://api.openai.com")
    .post("/v1/chat/completions")
    .reply(200, () => {
      // 模拟SSE格式的多块数据
      return [
        `data: ${JSON.stringify({
          id: "chatcmpl-123",
          object: "chat.completion.chunk",
          created: 1694268190,
          model: "gpt-3.5-turbo-0613",
          choices: [{ index: 0, delta: { role: "assistant", content: "I am" } }]
        })}\n\n`,
        `data: ${JSON.stringify({
          id: "chatcmpl-123",
          object: "chat.completion.chunk",
          created: 1694268190,
          model: "gpt-3.5-turbo-0613",
          choices: [{ index: 0, delta: { content: " actually not OpenAI." } }]
        })}\n\n`,
        `data: [DONE]\n\n`
      ].join('');
    }, {
      'Content-Type': 'text/event-stream',
      'Transfer-Encoding': 'chunked'
    });

  const response = await agent.get(`${baseRoute}/documents/${documentId}`);
  expect(response.status).toBe(200);
  expect(response.text).toBe("I am actually not OpenAI.");
});

方案2：直接Mock OpenAI SDK（更可靠） #

绕过HTTP拦截，直接mock openai.chat.completions.create方法，返回模拟的异步迭代器：

// Mock openaiService模块
jest.mock('../path/to/openaiService', () => {
  // 模拟流式响应的异步生成器
  const mockStream = async function*() {
    yield { choices: [{ delta: { role: "assistant", content: "I am" } }] };
    yield { choices: [{ delta: { content: " actually not OpenAI." } }] };
  };

  return {
    getCompletionStream: jest.fn().mockReturnValue(mockStream())
  };
});

it("should return stream the ai output if this is the first time fetching the document", async () => {
  const response = await agent.get(`${baseRoute}/documents/${documentId}`);
  expect(response.status).toBe(200);
  expect(response.text).toBe("I am actually not OpenAI.");
  // 验证方法被调用
  expect(require('../path/to/openaiService').getCompletionStream).toHaveBeenCalled();
});

方案3：监听流式响应的分块传输 #

如果需要验证每一块的传输过程，用supertest的on('data')事件监听：

it("should stream ai output chunks", (done) => {
  nock("https://api.openai.com")
    .post("/v1/chat/completions")
    .reply(200, function() {
      const chunks = [
        `data: ${JSON.stringify({ choices: [{ delta: { content: "I am" } }] })}\n\n`,
        `data: ${JSON.stringify({ choices: [{ delta: { content: " actually not" } }] })}\n\n`,
        `data: ${JSON.stringify({ choices: [{ delta: { content: " OpenAI." } }] })}\n\n`,
        `data: [DONE]\n\n`
      ];
      chunks.forEach(chunk => this.response.write(chunk));
      this.response.end();
    }, { 'Content-Type': 'text/event-stream' });

  agent.get(`${baseRoute}/documents/${documentId}`)
    .expect(200)
    .on('data', (chunk) => {
      // 验证每一块内容
      expect(chunk.toString()).toMatch(/I am|actually not|OpenAI/);
    })
    .on('end', done);
});

内容的提问来源于stack exchange，提问作者user16197389