From 1e9eedf38e83934bbc41fe6ddaefa23c2ed9f3ec Mon Sep 17 00:00:00 2001
From: Lahiru Maramba <llahiru@gmail.com>
Date: Wed, 29 May 2024 16:09:24 -0400
Subject: [PATCH] Add unit tests

---
 .../main/src/models/generative-model.test.ts  | 24 ++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/packages/main/src/models/generative-model.test.ts b/packages/main/src/models/generative-model.test.ts
index 5f4ce1c9..7de4715f 100644
--- a/packages/main/src/models/generative-model.test.ts
+++ b/packages/main/src/models/generative-model.test.ts
@@ -19,6 +19,7 @@ import { GenerativeModel } from "./generative-model";
 import * as sinonChai from "sinon-chai";
 import {
   FunctionCallingMode,
+  FunctionDeclarationSchemaType,
   HarmBlockThreshold,
   HarmCategory,
 } from "../../types";
@@ -50,7 +51,19 @@ describe("GenerativeModel", () => {
       "apiKey",
       {
         model: "my-model",
-        generationConfig: { temperature: 0 },
+        generationConfig: {
+          temperature: 0,
+          responseMimeType: "application/json",
+          responseSchema: {
+            type: FunctionDeclarationSchemaType.OBJECT,
+            properties: {
+              testField: {
+                type: FunctionDeclarationSchemaType.STRING,
+                properties: {},
+              },
+            },
+          },
+        },
         safetySettings: [
           {
             category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
@@ -68,6 +81,15 @@ describe("GenerativeModel", () => {
       },
     );
     expect(genModel.generationConfig?.temperature).to.equal(0);
+    expect(genModel.generationConfig?.responseMimeType).to.equal(
+      "application/json",
+    );
+    expect(genModel.generationConfig?.responseSchema.type).to.equal(
+      FunctionDeclarationSchemaType.OBJECT,
+    );
+    expect(
+      genModel.generationConfig?.responseSchema.properties.testField.type,
+    ).to.equal(FunctionDeclarationSchemaType.STRING);
     expect(genModel.safetySettings?.length).to.equal(1);
     expect(genModel.tools?.length).to.equal(1);
     expect(genModel.toolConfig?.functionCallingConfig.mode).to.equal(