Skip to content

Commit c4eabcb

Browse files
DellaBittaDellaBitta
committed
Fix some flaky generation tests
1 parent 5940794 commit c4eabcb

File tree

2 files changed

+21
-61
lines changed

2 files changed

+21
-61
lines changed

‎packages/ai/integration/chat.test.ts‎

Lines changed: 8 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -102,58 +102,29 @@ describe('Chat Session', () => {
102102
17, // "What is the capital of France?" + system instruction
103103
TOKEN_COUNT_DELTA + 2 // More variance for chat context
104104
);
105-
expect(response1.usageMetadata!.candidatesTokenCount).to.be.closeTo(
106-
8, // "Paris"
107-
TOKEN_COUNT_DELTA
108-
);
109-
expect(response1.usageMetadata!.totalTokenCount).to.be.closeTo(
110-
49, // "What is the capital of France?" + system instruction + "Paris"
111-
TOKEN_COUNT_DELTA + 3 // More variance for chat context
112-
);
113-
expect(response1.usageMetadata!.totalTokenCount).to.be.closeTo(
114-
49, // "What is the capital of France?" + system instruction + "Paris"
115-
TOKEN_COUNT_DELTA + 3 // More variance for chat context
116-
);
117-
105+
expect(response1.usageMetadata!.candidatesTokenCount).to.not.equal(0);
106+
expect(response1.usageMetadata!.totalTokenCount).to.not.equal(0);
118107
expect(response2.usageMetadata!.promptTokenCount).to.be.closeTo(
119108
32, // History + "And what about Italy?" + system instruction
120109
TOKEN_COUNT_DELTA + 5 // More variance for chat context with history
121110
);
122-
expect(response2.usageMetadata!.candidatesTokenCount).to.be.closeTo(
123-
8,
124-
TOKEN_COUNT_DELTA
125-
);
126-
expect(response2.usageMetadata!.totalTokenCount).to.be.closeTo(
127-
68,
128-
TOKEN_COUNT_DELTA + 2
129-
);
111+
expect(response2.usageMetadata!.candidatesTokenCount).to.not.equal(0);
112+
expect(response2.usageMetadata!.totalTokenCount).to.not.equal(0);
130113
} else if (model.model.includes('gemini-2.0-flash')) {
131114
expect(response1.usageMetadata).to.not.be.null;
132115
// Token counts can vary slightly in chat context
133116
expect(response1.usageMetadata!.promptTokenCount).to.be.closeTo(
134117
15, // "What is the capital of France?" + system instruction
135118
TOKEN_COUNT_DELTA + 2 // More variance for chat context
136119
);
137-
expect(response1.usageMetadata!.candidatesTokenCount).to.be.closeTo(
138-
8, // "Paris"
139-
TOKEN_COUNT_DELTA
140-
);
141-
expect(response1.usageMetadata!.totalTokenCount).to.be.closeTo(
142-
23, // "What is the capital of France?" + system instruction + "Paris"
143-
TOKEN_COUNT_DELTA + 3 // More variance for chat context
144-
);
120+
expect(response1.usageMetadata!.candidatesTokenCount).to.not.equal(0);
121+
expect(response1.usageMetadata!.totalTokenCount).to.not.equal(0);
145122
expect(response2.usageMetadata!.promptTokenCount).to.be.closeTo(
146123
28, // History + "And what about Italy?" + system instruction
147124
TOKEN_COUNT_DELTA + 5 // More variance for chat context with history
148125
);
149-
expect(response2.usageMetadata!.candidatesTokenCount).to.be.closeTo(
150-
8,
151-
TOKEN_COUNT_DELTA
152-
);
153-
expect(response2.usageMetadata!.totalTokenCount).to.be.closeTo(
154-
36,
155-
TOKEN_COUNT_DELTA
156-
);
126+
expect(response2.usageMetadata!.candidatesTokenCount).to.not.equal(0);
127+
expect(response2.usageMetadata!.totalTokenCount).to.not.equal(0);
157128
}
158129
});
159130
});

‎packages/ai/integration/generate-content.test.ts‎

Lines changed: 13 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -92,18 +92,9 @@ describe('Generate Content', function () {
9292
22,
9393
TOKEN_COUNT_DELTA
9494
);
95-
expect(response.usageMetadata!.candidatesTokenCount).to.be.closeTo(
96-
2,
97-
TOKEN_COUNT_DELTA
98-
);
99-
expect(response.usageMetadata!.thoughtsTokenCount).to.be.closeTo(
100-
30,
101-
TOKEN_COUNT_DELTA * 2
102-
);
103-
expect(response.usageMetadata!.totalTokenCount).to.be.closeTo(
104-
55,
105-
TOKEN_COUNT_DELTA * 2
106-
);
95+
expect(response.usageMetadata!.candidatesTokenCount).to.not.equal(0);
96+
expect(response.usageMetadata!.thoughtsTokenCount).to.not.equal(0);
97+
expect(response.usageMetadata!.totalTokenCount).to.not.equal(0);
10798
expect(response.usageMetadata!.promptTokensDetails).to.not.be.null;
10899
expect(response.usageMetadata!.promptTokensDetails!.length).to.equal(
109100
1
@@ -121,14 +112,8 @@ describe('Generate Content', function () {
121112
21,
122113
TOKEN_COUNT_DELTA
123114
);
124-
expect(response.usageMetadata!.candidatesTokenCount).to.be.closeTo(
125-
4,
126-
TOKEN_COUNT_DELTA
127-
);
128-
expect(response.usageMetadata!.totalTokenCount).to.be.closeTo(
129-
25,
130-
TOKEN_COUNT_DELTA * 2
131-
);
115+
expect(response.usageMetadata!.candidatesTokenCount).to.not.equal(0);
116+
expect(response.usageMetadata!.totalTokenCount).to.not.equal(0);
132117
expect(response.usageMetadata!.promptTokensDetails).to.not.be.null;
133118
expect(response.usageMetadata!.promptTokensDetails!.length).to.equal(
134119
1
@@ -149,7 +134,7 @@ describe('Generate Content', function () {
149134
).to.equal(Modality.TEXT);
150135
expect(
151136
response.usageMetadata!.candidatesTokensDetails![0].tokenCount
152-
).to.be.closeTo(4, TOKEN_COUNT_DELTA);
137+
).to.not.equal(0);
153138
}
154139
});
155140

@@ -230,8 +215,12 @@ describe('Generate Content', function () {
230215

231216
const usageMetadata = response.usageMetadata;
232217
expect(usageMetadata).to.exist;
233-
expect(usageMetadata?.toolUsePromptTokenCount).to.exist;
234-
expect(usageMetadata?.toolUsePromptTokenCount).to.be.greaterThan(0);
218+
// usageMetaData.toolUsePromptTokenCount does not exist in Gemini 2.0 flash responses.
219+
if (!model.model.includes('gemini-2.0-flash')) {
220+
221+
expect(usageMetadata?.toolUsePromptTokenCount).to.exist;
222+
expect(usageMetadata?.toolUsePromptTokenCount).to.be.greaterThan(0);
223+
}
235224
});
236225

237226
it('generateContent: url context and google search grounding', async () => {
@@ -288,7 +277,7 @@ describe('Generate Content', function () {
288277
});
289278

290279
const result = await model.generateContent(
291-
'Recommend 3 books for beginners to read to learn more about the latest advancements in Quantum Computing.'
280+
'Recommend 3 books for beginners to read to learn more about the latest advancements in Quantum Computing'
292281
);
293282
const response = result.response;
294283
const urlContextMetadata =

0 commit comments

Comments
 (0)