test(ai): update unit tests to use updated mocks

Author: russellwheatley

packages/ai/__tests__/generative-model.test.ts

@@ -19,7 +19,7 @@ import { type ReactNativeFirebase } from '@react-native-firebase/app';
 import { GenerativeModel } from '../lib/models/generative-model';
 import { AI, FunctionCallingMode } from '../lib/public-types';
 import * as request from '../lib/requests/request';
-import { getMockResponse } from './test-utils/mock-response';
+import { BackendName, getMockResponse } from './test-utils/mock-response';
 import { VertexAIBackend } from '../lib/backend';
 
 const fakeAI: AI = {
@@ -56,7 +56,10 @@ describe('GenerativeModel', () => {
     expect(genModel.tools?.length).toBe(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).toBe(FunctionCallingMode.NONE);
     expect(genModel.systemInstruction?.parts[0]!.text).toBe('be friendly');
-    const mockResponse = getMockResponse('unary-success-basic-reply-short.json');
+    const mockResponse = getMockResponse(
+      BackendName.VertexAI,
+      'unary-success-basic-reply-short.json',
+    );
     const makeRequestStub = jest
       .spyOn(request, 'makeRequest')
       .mockResolvedValue(mockResponse as Response);
@@ -78,7 +81,10 @@ describe('GenerativeModel', () => {
       systemInstruction: 'be friendly',
     });
     expect(genModel.systemInstruction?.parts[0]!.text).toBe('be friendly');
-    const mockResponse = getMockResponse('unary-success-basic-reply-short.json');
+    const mockResponse = getMockResponse(
+      BackendName.VertexAI,
+      'unary-success-basic-reply-short.json',
+    );
     const makeRequestStub = jest
       .spyOn(request, 'makeRequest')
       .mockResolvedValue(mockResponse as Response);
@@ -113,7 +119,10 @@ describe('GenerativeModel', () => {
     expect(genModel.tools?.length).toBe(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).toBe(FunctionCallingMode.NONE);
     expect(genModel.systemInstruction?.parts[0]!.text).toBe('be friendly');
-    const mockResponse = getMockResponse('unary-success-basic-reply-short.json');
+    const mockResponse = getMockResponse(
+      BackendName.VertexAI,
+      'unary-success-basic-reply-short.json',
+    );
     const makeRequestStub = jest
       .spyOn(request, 'makeRequest')
       .mockResolvedValue(mockResponse as Response);
@@ -178,7 +187,10 @@ describe('GenerativeModel', () => {
     expect(genModel.tools?.length).toBe(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).toBe(FunctionCallingMode.NONE);
     expect(genModel.systemInstruction?.parts[0]!.text).toBe('be friendly');
-    const mockResponse = getMockResponse('unary-success-basic-reply-short.json');
+    const mockResponse = getMockResponse(
+      BackendName.VertexAI,
+      'unary-success-basic-reply-short.json',
+    );
     const makeRequestStub = jest
       .spyOn(request, 'makeRequest')
       .mockResolvedValue(mockResponse as Response);
@@ -200,7 +212,10 @@ describe('GenerativeModel', () => {
       systemInstruction: 'be friendly',
     });
     expect(genModel.systemInstruction?.parts[0]!.text).toBe('be friendly');
-    const mockResponse = getMockResponse('unary-success-basic-reply-short.json');
+    const mockResponse = getMockResponse(
+      BackendName.VertexAI,
+      'unary-success-basic-reply-short.json',
+    );
     const makeRequestStub = jest
       .spyOn(request, 'makeRequest')
       .mockResolvedValue(mockResponse as Response);
@@ -226,7 +241,10 @@ describe('GenerativeModel', () => {
     expect(genModel.tools?.length).toBe(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).toBe(FunctionCallingMode.NONE);
     expect(genModel.systemInstruction?.parts[0]!.text).toBe('be friendly');
-    const mockResponse = getMockResponse('unary-success-basic-reply-short.json');
+    const mockResponse = getMockResponse(
+      BackendName.VertexAI,
+      'unary-success-basic-reply-short.json',
+    );
     const makeRequestStub = jest
       .spyOn(request, 'makeRequest')
       .mockResolvedValue(mockResponse as Response);
@@ -256,7 +274,7 @@ describe('GenerativeModel', () => {
 
   it('calls countTokens', async () => {
     const genModel = new GenerativeModel(fakeAI, { model: 'my-model' });
-    const mockResponse = getMockResponse('unary-success-total-tokens.json');
+    const mockResponse = getMockResponse(BackendName.VertexAI, 'unary-success-total-tokens.json');
     const makeRequestStub = jest
       .spyOn(request, 'makeRequest')
       .mockResolvedValue(mockResponse as Response);

packages/ai/__tests__/googleai-mappers.test.ts

@@ -41,7 +41,7 @@ import {
   PromptFeedback,
   SafetyRating,
 } from '../lib/public-types';
-import { getMockResponse } from './test-utils/mock-response';
+import { BackendName, getMockResponse } from './test-utils/mock-response';
 import { SpiedFunction } from 'jest-mock';
 
 const fakeModel = 'models/gemini-pro';
@@ -129,7 +129,7 @@ describe('Google AI Mappers', () => {
   describe('mapGenerateContentResponse', () => {
     it('should map a full Google AI response', async () => {
       const googleAIMockResponse: GoogleAIGenerateContentResponse = await (
-        getMockResponse('unary-success-citations.json') as Response
+        getMockResponse(BackendName.GoogleAI, 'unary-success-citations.json') as Response
       ).json();
       const mappedResponse = mapGenerateContentResponse(googleAIMockResponse);
 

packages/ai/__tests__/request.test.ts

@@ -20,7 +20,7 @@ import { ApiSettings } from '../lib/types/internal';
 import { DEFAULT_API_VERSION } from '../lib/constants';
 import { AIErrorCode } from '../lib/types';
 import { AIError } from '../lib/errors';
-import { getMockResponse } from './test-utils/mock-response';
+import { BackendName, getMockResponse } from './test-utils/mock-response';
 import { VertexAIBackend } from '../lib/backend';
 
 const fakeApiSettings: ApiSettings = {
@@ -350,7 +350,10 @@ describe('request methods', () => {
   });
 
   it('Network error, API not enabled', async () => {
-    const mockResponse = getMockResponse('unary-failure-firebasevertexai-api-not-enabled.json');
+    const mockResponse = getMockResponse(
+      BackendName.VertexAI,
+      'unary-failure-firebasevertexai-api-not-enabled.json',
+    );
     const fetchMock = jest.spyOn(globalThis, 'fetch').mockResolvedValue(mockResponse as Response);
     try {
       await makeRequest('models/model-name', Task.GENERATE_CONTENT, fakeApiSettings, false, '');

packages/ai/__tests__/stream-reader.test.ts

@@ -22,7 +22,11 @@ import {
   processStream,
 } from '../lib/requests/stream-reader';
 
-import { getChunkedStream, getMockResponseStreaming } from './test-utils/mock-response';
+import {
+  BackendName,
+  getChunkedStream,
+  getMockResponseStreaming,
+} from './test-utils/mock-response';
 import {
   BlockReason,
   FinishReason,
@@ -95,7 +99,10 @@ describe('stream-reader', () => {
     });
 
     it('streaming response - short', async () => {
-      const fakeResponse = getMockResponseStreaming('streaming-success-basic-reply-short.txt');
+      const fakeResponse = getMockResponseStreaming(
+        BackendName.VertexAI,
+        'streaming-success-basic-reply-short.txt',
+      );
       const result = processStream(fakeResponse as Response, fakeApiSettings);
       for await (const response of result.stream) {
         expect(response.text()).not.toBe('');
@@ -105,7 +112,10 @@ describe('stream-reader', () => {
     });
 
     it('streaming response - functioncall', async () => {
-      const fakeResponse = getMockResponseStreaming('streaming-success-function-call-short.txt');
+      const fakeResponse = getMockResponseStreaming(
+        BackendName.VertexAI,
+        'streaming-success-function-call-short.txt',
+      );
       const result = processStream(fakeResponse as Response, fakeApiSettings);
       for await (const response of result.stream) {
         expect(response.text()).toBe('');
@@ -127,7 +137,10 @@ describe('stream-reader', () => {
     });
 
     it('handles citations', async () => {
-      const fakeResponse = getMockResponseStreaming('streaming-success-citations.txt');
+      const fakeResponse = getMockResponseStreaming(
+        BackendName.VertexAI,
+        'streaming-success-citations.txt',
+      );
       const result = processStream(fakeResponse as Response, fakeApiSettings);
       const aggregatedResponse = await result.response;
       expect(aggregatedResponse.text()).toContain('Quantum mechanics is');
@@ -143,7 +156,10 @@ describe('stream-reader', () => {
     });
 
     it('removes empty text parts', async () => {
-      const fakeResponse = getMockResponseStreaming('streaming-success-empty-text-part.txt');
+      const fakeResponse = getMockResponseStreaming(
+        BackendName.VertexAI,
+        'streaming-success-empty-text-part.txt',
+      );
       const result = processStream(fakeResponse as Response, fakeApiSettings);
       const aggregatedResponse = await result.response;
       expect(aggregatedResponse.text()).toBe('1');