Skip to content

Commit

Permalink
test(firestore-multimodal-genai): get tests passing
Browse files Browse the repository at this point in the history
  • Loading branch information
@cabljac
cabljac committed Nov 28, 2024
1 parent 0353027 commit 038ca1d
Show file tree
Hide file tree
Showing 7 changed files with 110 additions and 44 deletions.
26 changes: 19 additions & 7 deletions firestore-multimodal-genai/functions/__tests__/google_ai/images.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}',
Expand All @@ -33,8 +34,8 @@ jest.mock('../../src/config', () => ({
provider: 'google-ai',
candidates: {
field: 'candidates',
count: 1,
shouldIncludeCandidatesField: false,
count: 2,
shouldIncludeCandidatesField: true,
},
},
}));
Expand Down Expand Up @@ -78,6 +79,15 @@ jest.mock('@google/generative-ai', () => {
],
},
},
{
content: {
parts: [
{
text: 'test response',
},
],
},
},
],
},
};
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand Down
25 changes: 18 additions & 7 deletions firestore-multimodal-genai/functions/__tests__/google_ai/index.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}',
Expand All @@ -32,8 +34,8 @@ jest.mock('../../src/config', () => ({
provider: 'google-ai',
candidates: {
field: 'candidates',
count: 1,
shouldIncludeCandidatesField: false,
count: 2,
shouldIncludeCandidatesField: true,
},
},
}));
Expand Down Expand Up @@ -66,6 +68,15 @@ jest.mock('@google/generative-ai', () => {
],
},
},
{
content: {
parts: [
{
text: 'test response',
},
],
},
},
],
},
};
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand Down
43 changes: 32 additions & 11 deletions firestore-multimodal-genai/functions/__tests__/util.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down
24 changes: 17 additions & 7 deletions firestore-multimodal-genai/functions/__tests__/vertex_ai/images.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}',
Expand All @@ -33,8 +34,8 @@ jest.mock('../../src/config', () => ({
provider: 'vertex-ai',
candidates: {
field: 'candidates',
count: 1,
shouldIncludeCandidatesField: false,
count: 2,
shouldIncludeCandidatesField: true,
},
},
}));
Expand Down Expand Up @@ -75,6 +76,15 @@ jest.mock('@google-cloud/vertexai', () => {
],
},
},
{
content: {
parts: [
{
text: 'test response',
},
],
},
},
],
},
};
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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
Expand Down
24 changes: 17 additions & 7 deletions firestore-multimodal-genai/functions/__tests__/vertex_ai/index.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}',
Expand All @@ -32,8 +33,8 @@ jest.mock('../../src/config', () => ({
provider: 'vertex-ai',
candidates: {
field: 'candidates',
count: 1,
shouldIncludeCandidatesField: false,
count: 2,
shouldIncludeCandidatesField: true,
},
},
}));
Expand Down Expand Up @@ -67,6 +68,15 @@ jest.mock('@google-cloud/vertexai', () => {
],
},
},
{
content: {
parts: [
{
text: 'test response',
},
],
},
},
],
},
};
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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
Expand Down
9 changes: 5 additions & 4 deletions firestore-multimodal-genai/functions/src/__mocks__/config.ts
View file
Original file line number Diff line number Diff line change
@@ -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',
Expand All @@ -18,7 +19,7 @@ export default {
imageField: 'image',
candidates: {
field: 'candidates',
count: 1,
shouldIncludeCandidatesField: false,
count: 5,
shouldIncludeCandidatesField: true,
},
};
3 changes: 2 additions & 1 deletion firestore-multimodal-genai/functions/src/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down

0 comments on commit 038ca1d

Please sign in to comment.