backend/node_modules/@google/generative-ai/src/methods/chat-session.test.ts

/**
 * @license
 * Copyright 2024 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { expect, use } from "chai";
import { match, restore, stub, useFakeTimers } from "sinon";
import * as sinonChai from "sinon-chai";
import * as chaiAsPromised from "chai-as-promised";
import * as generateContentMethods from "./generate-content";
import { GenerateContentStreamResult } from "../../types";
import { ChatSession } from "./chat-session";

use(sinonChai);
use(chaiAsPromised);

describe("ChatSession", () => {
  afterEach(() => {
    restore();
  });
  describe("sendMessage()", () => {
    it("generateContent errors should be catchable", async () => {
      const generateContentStub = stub(
        generateContentMethods,
        "generateContent",
      ).rejects("generateContent failed");
      const chatSession = new ChatSession("MY_API_KEY", "a-model");
      await expect(chatSession.sendMessage("hello")).to.be.rejected;
      expect(generateContentStub).to.be.calledWith(
        "MY_API_KEY",
        "a-model",
        match.any,
      );
    });
  });
  describe("sendMessageStream()", () => {
    it("generateContentStream errors should be catchable", async () => {
      const clock = useFakeTimers();
      const consoleStub = stub(console, "error");
      const generateContentStreamStub = stub(
        generateContentMethods,
        "generateContentStream",
      ).rejects("generateContentStream failed");
      const chatSession = new ChatSession("MY_API_KEY", "a-model");
      await expect(chatSession.sendMessageStream("hello")).to.be.rejected;
      expect(generateContentStreamStub).to.be.calledWith(
        "MY_API_KEY",
        "a-model",
        match.any,
      );
      await clock.runAllAsync();
      expect(consoleStub).to.not.be.called;
      clock.restore();
    });
    it("downstream sendPromise errors should log but not throw", async () => {
      const clock = useFakeTimers();
      const consoleStub = stub(console, "error");
      // make candidates error when read as an array
      const generateContentStreamStub = stub(
        generateContentMethods,
        "generateContentStream",
      ).resolves({ candidates: 3 } as unknown as GenerateContentStreamResult);
      const chatSession = new ChatSession("MY_API_KEY", "a-model");
      await chatSession.sendMessageStream("hello");
      expect(generateContentStreamStub).to.be.calledWith(
        "MY_API_KEY",
        "a-model",
        match.any,
      );
      await clock.runAllAsync();
      expect(consoleStub.args[0][0].toString()).to.include(
        "properties of undefined",
      );
      clock.restore();
    });
  });
});