Skip to content

Commit 32715cf

Browse files
dlarocquedlarocque
committed
Add unit tests to model classes
1 parent cc21c93 commit 32715cf

File tree

6 files changed

+381
-18
lines changed

6 files changed

+381
-18
lines changed

‎packages/vertexai/src/methods/chat-session.test.ts‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,64 @@ describe('ChatSession', () => {
5252
match.any
5353
);
5454
});
55+
it('singleRequestOptions overrides requestOptions', async () => {
56+
const generateContentStub = stub(
57+
generateContentMethods,
58+
'generateContent'
59+
).rejects('generateContent failed'); // not important
60+
const requestOptions = {
61+
timeout: 1000
62+
};
63+
const singleRequestOptions = {
64+
timeout: 2000
65+
};
66+
const chatSession = new ChatSession(
67+
fakeApiSettings,
68+
'a-model',
69+
undefined,
70+
requestOptions
71+
);
72+
await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
73+
.rejected;
74+
expect(generateContentStub).to.be.calledWith(
75+
fakeApiSettings,
76+
'a-model',
77+
match.any,
78+
match({
79+
timeout: singleRequestOptions.timeout
80+
})
81+
);
82+
});
83+
it('singleRequestOptions is merged with requestOptions', async () => {
84+
const generateContentStub = stub(
85+
generateContentMethods,
86+
'generateContent'
87+
).rejects('generateContent failed'); // not important
88+
const abortController = new AbortController();
89+
const requestOptions = {
90+
timeout: 1000
91+
};
92+
const singleRequestOptions = {
93+
signal: abortController.signal
94+
};
95+
const chatSession = new ChatSession(
96+
fakeApiSettings,
97+
'a-model',
98+
undefined,
99+
requestOptions
100+
);
101+
await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
102+
.rejected;
103+
expect(generateContentStub).to.be.calledWith(
104+
fakeApiSettings,
105+
'a-model',
106+
match.any,
107+
match({
108+
timeout: requestOptions.timeout,
109+
signal: singleRequestOptions.signal
110+
})
111+
);
112+
});
55113
});
56114
describe('sendMessageStream()', () => {
57115
it('generateContentStream errors should be catchable', async () => {
@@ -94,5 +152,63 @@ describe('ChatSession', () => {
94152
);
95153
clock.restore();
96154
});
155+
it('singleRequestOptions overrides requestOptions', async () => {
156+
const generateContentStub = stub(
157+
generateContentMethods,
158+
'generateContent'
159+
).rejects('generateContent failed'); // not important
160+
const requestOptions = {
161+
timeout: 1000
162+
};
163+
const singleRequestOptions = {
164+
timeout: 2000
165+
};
166+
const chatSession = new ChatSession(
167+
fakeApiSettings,
168+
'a-model',
169+
undefined,
170+
requestOptions
171+
);
172+
await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
173+
.rejected;
174+
expect(generateContentStub).to.be.calledWith(
175+
fakeApiSettings,
176+
'a-model',
177+
match.any,
178+
match({
179+
timeout: singleRequestOptions.timeout
180+
})
181+
);
182+
});
183+
it('singleRequestOptions is merged with requestOptions', async () => {
184+
const generateContentStub = stub(
185+
generateContentMethods,
186+
'generateContent'
187+
).rejects('generateContent failed'); // not important
188+
const abortController = new AbortController();
189+
const requestOptions = {
190+
timeout: 1000
191+
};
192+
const singleRequestOptions = {
193+
signal: abortController.signal
194+
};
195+
const chatSession = new ChatSession(
196+
fakeApiSettings,
197+
'a-model',
198+
undefined,
199+
requestOptions
200+
);
201+
await expect(chatSession.sendMessage('hello', singleRequestOptions)).to.be
202+
.rejected;
203+
expect(generateContentStub).to.be.calledWith(
204+
fakeApiSettings,
205+
'a-model',
206+
match.any,
207+
match({
208+
timeout: requestOptions.timeout,
209+
signal: singleRequestOptions.signal
210+
})
211+
);
212+
});
97213
});
98214
});

‎packages/vertexai/src/methods/chat-session.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,8 @@ export class ChatSession {
9999
generateContentRequest,
100100
// Merge requestOptions
101101
{
102-
...singleRequestOptions,
103-
...this.requestOptions
102+
...this.requestOptions,
103+
...singleRequestOptions
104104
}
105105
)
106106
)
@@ -155,8 +155,8 @@ export class ChatSession {
155155
generateContentRequest,
156156
// Merge requestOptions
157157
{
158-
...singleRequestOptions,
159-
...this.requestOptions
158+
...this.requestOptions,
159+
...singleRequestOptions
160160
}
161161
);
162162

‎packages/vertexai/src/models/generative-model.test.ts‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,18 @@
1515
* limitations under the License.
1616
*/
1717
import { use, expect } from 'chai';
18+
import chaiAsPromised from 'chai-as-promised';
1819
import { GenerativeModel } from './generative-model';
1920
import { FunctionCallingMode, VertexAI } from '../public-types';
2021
import * as request from '../requests/request';
2122
import { match, restore, stub } from 'sinon';
2223
import { getMockResponse } from '../../test-utils/mock-response';
2324
import sinonChai from 'sinon-chai';
25+
import * as generateContentMethods from '../methods/generate-content';
26+
import * as countTokens from '../methods/count-tokens';
2427

2528
use(sinonChai);
29+
use(chaiAsPromised);
2630

2731
const fakeVertexAI: VertexAI = {
2832
app: {
@@ -38,6 +42,9 @@ const fakeVertexAI: VertexAI = {
3842
};
3943

4044
describe('GenerativeModel', () => {
45+
afterEach(() => {
46+
restore();
47+
});
4148
it('passes params through to generateContent', async () => {
4249
const genModel = new GenerativeModel(fakeVertexAI, {
4350
model: 'my-model',
@@ -165,6 +172,62 @@ describe('GenerativeModel', () => {
165172
);
166173
restore();
167174
});
175+
it('generateContent singleRequestOptions overrides requestOptions', async () => {
176+
const generateContentStub = stub(
177+
generateContentMethods,
178+
'generateContent'
179+
).rejects('generateContent failed'); // not important
180+
const requestOptions = {
181+
timeout: 1000
182+
};
183+
const singleRequestOptions = {
184+
timeout: 2000
185+
};
186+
const genModel = new GenerativeModel(
187+
fakeVertexAI,
188+
{ model: 'my-model' },
189+
requestOptions
190+
);
191+
await expect(genModel.generateContent('hello', singleRequestOptions)).to.be
192+
.rejected;
193+
expect(generateContentStub).to.be.calledWith(
194+
match.any,
195+
match.any,
196+
match.any,
197+
match({
198+
timeout: singleRequestOptions.timeout
199+
})
200+
);
201+
});
202+
it('generateContent singleRequestOptions is merged with requestOptions', async () => {
203+
const generateContentStub = stub(
204+
generateContentMethods,
205+
'generateContent'
206+
).rejects('generateContent failed'); // not important
207+
const abortController = new AbortController();
208+
const requestOptions = {
209+
timeout: 1000
210+
};
211+
const singleRequestOptions = {
212+
signal: abortController.signal
213+
};
214+
const genModel = new GenerativeModel(
215+
fakeVertexAI,
216+
{ model: 'my-model' },
217+
requestOptions
218+
);
219+
await expect(genModel.generateContent('hello', singleRequestOptions)).to.be
220+
.rejected;
221+
expect(generateContentStub).to.be.calledWith(
222+
match.any,
223+
match.any,
224+
match.any,
225+
match({
226+
timeout: requestOptions.timeout,
227+
signal: singleRequestOptions.signal
228+
})
229+
);
230+
});
168231
it('passes params through to chat.sendMessage', async () => {
169232
const genModel = new GenerativeModel(fakeVertexAI, {
170233
model: 'my-model',
@@ -314,4 +377,57 @@ describe('GenerativeModel', () => {
314377
);
315378
restore();
316379
});
380+
it('countTokens singleRequestOptions overrides requestOptions', async () => {
381+
const countTokensStub = stub(countTokens, 'countTokens').rejects(
382+
'countTokens failed'
383+
);
384+
const requestOptions = {
385+
timeout: 1000
386+
};
387+
const singleRequestOptions = {
388+
timeout: 2000
389+
};
390+
const genModel = new GenerativeModel(
391+
fakeVertexAI,
392+
{ model: 'my-model' },
393+
requestOptions
394+
);
395+
await expect(genModel.countTokens('hello', singleRequestOptions)).to.be.rejected;
396+
expect(countTokensStub).to.be.calledWith(
397+
match.any,
398+
match.any,
399+
match.any,
400+
match({
401+
timeout: singleRequestOptions.timeout
402+
})
403+
);
404+
});
405+
it('countTokens singleRequestOptions is merged with requestOptions', async () => {
406+
const countTokensStub = stub(countTokens, 'countTokens').rejects(
407+
'countTokens failed'
408+
);
409+
const abortController = new AbortController();
410+
const requestOptions = {
411+
timeout: 1000
412+
};
413+
const singleRequestOptions = {
414+
signal: abortController.signal
415+
};
416+
const genModel = new GenerativeModel(
417+
fakeVertexAI,
418+
{ model: 'my-model' },
419+
requestOptions
420+
);
421+
await expect(genModel.countTokens('hello', singleRequestOptions)).to.be
422+
.rejected;
423+
expect(countTokensStub).to.be.calledWith(
424+
match.any,
425+
match.any,
426+
match.any,
427+
match({
428+
timeout: requestOptions.timeout,
429+
signal: singleRequestOptions.signal
430+
})
431+
);
432+
});
317433
});

‎packages/vertexai/src/models/generative-model.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,8 @@ export class GenerativeModel extends VertexAIModel {
9595
},
9696
// Merge request options
9797
{
98-
...singleRequestOptions,
99-
...this.requestOptions
98+
...this.requestOptions,
99+
...singleRequestOptions
100100
}
101101
);
102102
}
@@ -125,8 +125,8 @@ export class GenerativeModel extends VertexAIModel {
125125
},
126126
// Merge request options
127127
{
128-
...singleRequestOptions,
129-
...this.requestOptions
128+
...this.requestOptions,
129+
...singleRequestOptions
130130
}
131131
);
132132
}
@@ -170,8 +170,8 @@ export class GenerativeModel extends VertexAIModel {
170170
formattedParams,
171171
// Merge request options
172172
{
173-
...singleRequestOptions,
174-
...this.requestOptions
173+
...this.requestOptions,
174+
...singleRequestOptions
175175
}
176176
);
177177
}

0 commit comments

Comments
 (0)