From 038ca1d809b2dc5fd99eb76320acd2d8579f4686 Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Thu, 28 Nov 2024 15:23:02 +0000 Subject: [PATCH] test(firestore-multimodal-genai): get tests passing --- .../__tests__/google_ai/images.test.ts | 26 ++++++++--- .../__tests__/google_ai/index.test.ts | 25 ++++++++--- .../functions/__tests__/util.ts | 43 ++++++++++++++----- .../__tests__/vertex_ai/images.test.ts | 24 ++++++++--- .../__tests__/vertex_ai/index.test.ts | 24 ++++++++--- .../functions/src/__mocks__/config.ts | 9 ++-- .../functions/src/index.ts | 3 +- 7 files changed, 110 insertions(+), 44 deletions(-) diff --git a/firestore-multimodal-genai/functions/__tests__/google_ai/images.test.ts b/firestore-multimodal-genai/functions/__tests__/google_ai/images.test.ts index 28660c6d..cdc36494 100644 --- a/firestore-multimodal-genai/functions/__tests__/google_ai/images.test.ts +++ b/firestore-multimodal-genai/functions/__tests__/google_ai/images.test.ts @@ -16,12 +16,13 @@ process.env.FIRESTORE_EMULATOR_HOST = '127.0.0.1:8080'; jest.mock('../../src/config', () => ({ default: { googleAi: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', apiKey: 'test-api-key', }, vertex: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', }, + model: 'gemini-1.0-pro', collectionName: 'generate', location: 'us-central1', prompt: '{{ instruction }}', @@ -33,8 +34,8 @@ jest.mock('../../src/config', () => ({ provider: 'google-ai', candidates: { field: 'candidates', - count: 1, - shouldIncludeCandidatesField: false, + count: 2, + shouldIncludeCandidatesField: true, }, }, })); @@ -78,6 +79,15 @@ jest.mock('@google/generative-ai', () => { ], }, }, + { + content: { + parts: [ + { + text: 'test response', + }, + ], + }, + }, ], }, }; @@ -115,7 +125,7 @@ const wrappedGenerateMessage = fft.wrap( const firestoreObserver = jest.fn((_x: any) => {}); let collectionName: string; -describe('generateMessage', () => { +describe('generateMessage SDK directly', () => { let unsubscribe: (() => void) | undefined; // clear firestore @@ -249,7 +259,7 @@ describe('generateMessage', () => { call[0].docs[0].data() ); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); expect(mockGetClient).toHaveBeenCalledTimes(1); expect(mockGetClient).toHaveBeenCalledWith(config.googleAi.apiKey); @@ -313,7 +323,9 @@ describe('generateMessage', () => { return call[0].docs[0].data(); }); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + console.info('!!!', firestoreCallData); + + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); // // verify SDK is called with expected arguments // we expect the mock API to be called once diff --git a/firestore-multimodal-genai/functions/__tests__/google_ai/index.test.ts b/firestore-multimodal-genai/functions/__tests__/google_ai/index.test.ts index 4eb0e53b..5e14309a 100644 --- a/firestore-multimodal-genai/functions/__tests__/google_ai/index.test.ts +++ b/firestore-multimodal-genai/functions/__tests__/google_ai/index.test.ts @@ -16,12 +16,14 @@ process.env.FIRESTORE_EMULATOR_HOST = '127.0.0.1:8080'; jest.mock('../../src/config', () => ({ default: { googleAi: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', apiKey: 'test-api-key', }, vertex: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', }, + model: 'gemini-1.0-pro', + collectionName: 'generate', location: 'us-central1', prompt: '{{ instruction }}', @@ -32,8 +34,8 @@ jest.mock('../../src/config', () => ({ provider: 'google-ai', candidates: { field: 'candidates', - count: 1, - shouldIncludeCandidatesField: false, + count: 2, + shouldIncludeCandidatesField: true, }, }, })); @@ -66,6 +68,15 @@ jest.mock('@google/generative-ai', () => { ], }, }, + { + content: { + parts: [ + { + text: 'test response', + }, + ], + }, + }, ], }, }; @@ -103,7 +114,7 @@ const wrappedGenerateMessage = fft.wrap( const firestoreObserver = jest.fn((_x: any) => {}); let collectionName: string; -describe('generateMessage', () => { +describe('generateMessage SDK directly', () => { let unsubscribe: (() => void) | undefined; // clear firestore @@ -234,7 +245,7 @@ describe('generateMessage', () => { call[0].docs[0].data() ); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); expect(mockGetClient).toHaveBeenCalledTimes(1); expect(mockGetClient).toHaveBeenCalledWith(config.googleAi.apiKey); @@ -289,7 +300,7 @@ describe('generateMessage', () => { return call[0].docs[0].data(); }); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); // // verify SDK is called with expected arguments // we expect the mock API to be called once diff --git a/firestore-multimodal-genai/functions/__tests__/util.ts b/firestore-multimodal-genai/functions/__tests__/util.ts index 36db494f..fce2e471 100644 --- a/firestore-multimodal-genai/functions/__tests__/util.ts +++ b/firestore-multimodal-genai/functions/__tests__/util.ts @@ -3,7 +3,8 @@ import {Timestamp} from 'firebase-admin/firestore'; export const expectToProcessCorrectly = ( firestoreCallData: any[], message: any, - mockResponse = 'test response' + mockResponse = 'test response', + candidateCount?: number ) => { expect(firestoreCallData[0]).toEqual({ ...message, @@ -22,16 +23,36 @@ export const expectToProcessCorrectly = ( firestoreCallData[1].status.updateTime ); - expect(firestoreCallData[2]).toEqual({ - ...message, - output: mockResponse, - status: { - state: 'COMPLETED', - startTime: expect.any(Timestamp), - updateTime: expect.any(Timestamp), - completeTime: expect.any(Timestamp), - }, - }); + const expectedCandidates = + candidateCount !== undefined + ? Array(candidateCount).fill(mockResponse) + : []; + + const expectedCompleteData = + expectedCandidates.length > 0 + ? { + ...message, + output: mockResponse, + candidates: expectedCandidates, + status: { + state: 'COMPLETED', + startTime: expect.any(Timestamp), + updateTime: expect.any(Timestamp), + completeTime: expect.any(Timestamp), + }, + } + : { + ...message, + output: mockResponse, + status: { + state: 'COMPLETED', + startTime: expect.any(Timestamp), + updateTime: expect.any(Timestamp), + completeTime: expect.any(Timestamp), + }, + }; + + expect(firestoreCallData[2]).toEqual(expectedCompleteData); expect(firestoreCallData[2].status.startTime).toEqual( firestoreCallData[1].status.startTime diff --git a/firestore-multimodal-genai/functions/__tests__/vertex_ai/images.test.ts b/firestore-multimodal-genai/functions/__tests__/vertex_ai/images.test.ts index bc07433e..fde03d03 100644 --- a/firestore-multimodal-genai/functions/__tests__/vertex_ai/images.test.ts +++ b/firestore-multimodal-genai/functions/__tests__/vertex_ai/images.test.ts @@ -17,11 +17,12 @@ process.env.FIRESTORE_EMULATOR_HOST = '127.0.0.1:8080'; jest.mock('../../src/config', () => ({ default: { googleAi: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', }, vertex: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', }, + model: 'gemini-1.0-pro', collectionName: 'generate', location: 'us-central1', prompt: '{{ instruction }}', @@ -33,8 +34,8 @@ jest.mock('../../src/config', () => ({ provider: 'vertex-ai', candidates: { field: 'candidates', - count: 1, - shouldIncludeCandidatesField: false, + count: 2, + shouldIncludeCandidatesField: true, }, }, })); @@ -75,6 +76,15 @@ jest.mock('@google-cloud/vertexai', () => { ], }, }, + { + content: { + parts: [ + { + text: 'test response', + }, + ], + }, + }, ], }, }; @@ -113,7 +123,7 @@ const wrappedGenerateMessage = fft.wrap( const firestoreObserver = jest.fn((_x: any) => {}); let collectionName: string; -describe('generateMessage', () => { +describe('generateMessage SDK directly', () => { let unsubscribe: (() => void) | undefined; // clear firestore @@ -243,7 +253,7 @@ describe('generateMessage', () => { call[0].docs[0].data() ); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); expect(mockGetClient).toHaveBeenCalledTimes(1); @@ -302,7 +312,7 @@ describe('generateMessage', () => { return call[0].docs[0].data(); }); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); // // verify SDK is called with expected arguments // we expect the mock API to be called once diff --git a/firestore-multimodal-genai/functions/__tests__/vertex_ai/index.test.ts b/firestore-multimodal-genai/functions/__tests__/vertex_ai/index.test.ts index a7992cd8..1f6fd1fa 100644 --- a/firestore-multimodal-genai/functions/__tests__/vertex_ai/index.test.ts +++ b/firestore-multimodal-genai/functions/__tests__/vertex_ai/index.test.ts @@ -16,11 +16,12 @@ process.env.FIRESTORE_EMULATOR_HOST = '127.0.0.1:8080'; jest.mock('../../src/config', () => ({ default: { googleAi: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', }, vertex: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', }, + model: 'gemini-1.0-pro', collectionName: 'generate', location: 'us-central1', prompt: '{{ instruction }}', @@ -32,8 +33,8 @@ jest.mock('../../src/config', () => ({ provider: 'vertex-ai', candidates: { field: 'candidates', - count: 1, - shouldIncludeCandidatesField: false, + count: 2, + shouldIncludeCandidatesField: true, }, }, })); @@ -67,6 +68,15 @@ jest.mock('@google-cloud/vertexai', () => { ], }, }, + { + content: { + parts: [ + { + text: 'test response', + }, + ], + }, + }, ], }, }; @@ -105,7 +115,7 @@ const wrappedGenerateMessage = fft.wrap( const firestoreObserver = jest.fn((_x: any) => {}); let collectionName: string; -describe('generateMessage', () => { +describe('generateMessage SDK directly', () => { let unsubscribe: (() => void) | undefined; // clear firestore @@ -230,7 +240,7 @@ describe('generateMessage', () => { call[0].docs[0].data() ); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); expect(mockGetClient).toHaveBeenCalledTimes(1); @@ -283,7 +293,7 @@ describe('generateMessage', () => { return call[0].docs[0].data(); }); - expectToProcessCorrectly(firestoreCallData, message, 'test response'); + expectToProcessCorrectly(firestoreCallData, message, 'test response', 2); // // verify SDK is called with expected arguments // we expect the mock API to be called once diff --git a/firestore-multimodal-genai/functions/src/__mocks__/config.ts b/firestore-multimodal-genai/functions/src/__mocks__/config.ts index f139e54b..5fc4c595 100644 --- a/firestore-multimodal-genai/functions/src/__mocks__/config.ts +++ b/firestore-multimodal-genai/functions/src/__mocks__/config.ts @@ -1,11 +1,12 @@ export default { vertex: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', }, googleAi: { - model: 'gemini-pro', + model: 'gemini-1.0-pro', apiKey: 'test-api-key', }, + model: 'gemini-1.0-pro', location: 'us-central1', projectId: 'text-project-id', instanceId: 'text-instance-id', @@ -18,7 +19,7 @@ export default { imageField: 'image', candidates: { field: 'candidates', - count: 1, - shouldIncludeCandidatesField: false, + count: 5, + shouldIncludeCandidatesField: true, }, }; diff --git a/firestore-multimodal-genai/functions/src/index.ts b/firestore-multimodal-genai/functions/src/index.ts index abfe7fe8..8f497c72 100644 --- a/firestore-multimodal-genai/functions/src/index.ts +++ b/firestore-multimodal-genai/functions/src/index.ts @@ -118,7 +118,8 @@ export const generateText = functions.firestore ...metadata, [responseField]: result.candidates[0], [config.candidates.field]: result.candidates, - 'status.error': null, + 'status.state': 'COMPLETED', + // 'status.error': null, }); } return ref.update({