Skip to content

Commit e7312df

Browse files
svozzasvozza
committed
use it.each for testing fucntion return type tests
1 parent 823d4b5 commit e7312df

File tree

1 file changed

+124
-167
lines changed

1 file changed

+124
-167
lines changed

‎packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts

Lines changed: 124 additions & 167 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import context from '@aws-lambda-powertools/testing-utils/context';
22
import type { Context } from 'aws-lambda';
33
import { beforeEach, describe, expect, it, vi } from 'vitest';
4-
import type { BedrockAgentFunctionEvent, Parameter } from '../../src/types';
54
import { BedrockAgentFunctionResolver } from '../../src/bedrock-agent-function';
5+
import type { BedrockAgentFunctionEvent, Parameter } from '../../src/types';
66

77
function createEvent(functionName: string, parameters?: Parameter[]) {
88
return {
@@ -71,6 +71,8 @@ describe('Class: BedrockAgentFunctionResolver', () => {
7171
definition: 'Does nothing',
7272
});
7373

74+
// Act
75+
7476
await app.resolve(createEvent('noop', []), context);
7577

7678
// Assess
@@ -90,15 +92,14 @@ describe('Class: BedrockAgentFunctionResolver', () => {
9092
});
9193
}
9294

93-
app.tool(async (params) => {
94-
return Number(params.a) + Number(params);
95-
}, {
96-
name: 'mult',
97-
definition: 'Multiplies two numbers',
98-
});
99-
100-
expect(console.warn).toHaveBeenLastCalledWith(
101-
'The maximum number of tools that can be registered is 5. Tool mult will not be registered.'
95+
app.tool(
96+
async (params) => {
97+
return Number(params.a) + Number(params);
98+
},
99+
{
100+
name: 'mult',
101+
definition: 'Multiplies two numbers',
102+
}
102103
);
103104

104105
const event = createEvent('mult', [
@@ -114,8 +115,13 @@ describe('Class: BedrockAgentFunctionResolver', () => {
114115
},
115116
]);
116117

118+
// Act
117119
const actual = await app.resolve(event, context);
118120

121+
// Assess
122+
expect(console.warn).toHaveBeenLastCalledWith(
123+
'The maximum number of tools that can be registered is 5. Tool mult will not be registered.'
124+
);
119125
expect(actual.response.function).toEqual('mult');
120126
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
121127
'Error: tool has not been registered in handler.'
@@ -170,27 +176,6 @@ describe('Class: BedrockAgentFunctionResolver', () => {
170176
expect(
171177
multiplyResult.response.functionResponse.responseBody.TEXT.body
172178
).toEqual('20');
173-
174-
app.tool(
175-
async (params) => {
176-
return Number(params.a) / Number(params.b);
177-
},
178-
{
179-
name: 'math',
180-
definition: 'Divides two numbers',
181-
}
182-
);
183-
184-
const divideResult = await app.resolve(event, context);
185-
expect(divideResult.response.function).toEqual('math');
186-
expect(
187-
divideResult.response.functionResponse.responseBody.TEXT.body
188-
).toEqual('5');
189-
190-
expect(console.warn).toHaveBeenCalledTimes(3);
191-
expect(console.warn).toHaveBeenCalledWith(
192-
'Tool math already registered. Overwriting with new definition.'
193-
);
194179
});
195180

196181
it('can be invoked using the decorator pattern', async () => {
@@ -216,14 +201,6 @@ describe('Class: BedrockAgentFunctionResolver', () => {
216201

217202
const lambda = new Lambda();
218203

219-
const helloEvent = createEvent('hello');
220-
221-
const helloResult = await lambda.handler(helloEvent, context);
222-
expect(helloResult.response.function).toEqual('hello');
223-
expect(
224-
helloResult.response.functionResponse.responseBody.TEXT.body
225-
).toEqual('Hello, world!');
226-
227204
const addEvent = createEvent('add', [
228205
{
229206
name: 'a',
@@ -237,65 +214,43 @@ describe('Class: BedrockAgentFunctionResolver', () => {
237214
},
238215
]);
239216

240-
const addResult = await lambda.handler(addEvent, context);
241-
expect(addResult.response.function).toEqual('add');
242-
expect(addResult.response.functionResponse.responseBody.TEXT.body).toEqual(
243-
'3'
244-
);
245-
});
246-
247-
it('handles functions that return different primitive types', async () => {
248-
// Prepare
249-
const app = new BedrockAgentFunctionResolver();
250-
251-
app.tool(async () => 'Hello, world!', {
252-
name: 'string-tool',
253-
definition: 'Returns a string',
254-
});
217+
// Act
218+
const actual = await lambda.handler(addEvent, context);
255219

256-
app.tool(async () => 42, {
257-
name: 'number-tool',
258-
definition: 'Returns a number',
259-
});
260-
261-
app.tool(async () => true, {
262-
name: 'boolean-tool',
263-
definition: 'Returns a boolean',
264-
});
265-
266-
const stringResult = await app.resolve(createEvent('string-tool'), context);
267-
expect(stringResult.response.function).toEqual('string-tool');
268-
expect(
269-
stringResult.response.functionResponse.responseBody.TEXT.body
270-
).toEqual('Hello, world!');
271-
272-
const numberResult = await app.resolve(createEvent('number-tool'), context);
273-
expect(numberResult.response.function).toEqual('number-tool');
274-
expect(
275-
numberResult.response.functionResponse.responseBody.TEXT.body
276-
).toEqual('42');
277-
278-
const booleanResult = await app.resolve(
279-
createEvent('boolean-tool'),
280-
context
220+
// Assess
221+
expect(actual.response.function).toEqual('add');
222+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
223+
'3'
281224
);
282-
expect(booleanResult.response.function).toEqual('boolean-tool');
283-
expect(
284-
booleanResult.response.functionResponse.responseBody.TEXT.body
285-
).toEqual('true');
286225
});
287226

288-
it('handles functions that return complex types (array and object)', async () => {
289-
// Prepare
290-
const app = new BedrockAgentFunctionResolver();
291-
292-
app.tool(async () => [1, 'two', false, null], {
293-
name: 'array-tool',
294-
definition: 'Returns an array',
295-
});
296-
297-
app.tool(
298-
async () => ({
227+
it.each([
228+
{
229+
toolFunction: async () => 'Hello, world',
230+
toolParams: {
231+
name: 'string',
232+
definition: 'Returns string',
233+
},
234+
expected: 'Hello, world',
235+
},
236+
{
237+
toolFunction: async () => 42,
238+
toolParams: {
239+
name: 'number',
240+
definition: 'Returns number',
241+
},
242+
expected: '42',
243+
},
244+
{
245+
toolFunction: async () => true,
246+
toolParams: {
247+
name: 'boolean',
248+
definition: 'Returns boolean',
249+
},
250+
expected: 'true',
251+
},
252+
{
253+
toolFunction: async () => ({
299254
name: 'John Doe',
300255
age: 30,
301256
isActive: true,
@@ -304,70 +259,73 @@ describe('Class: BedrockAgentFunctionResolver', () => {
304259
city: 'Anytown',
305260
},
306261
}),
307-
{
308-
name: 'object-tool',
262+
toolParams: {
263+
name: 'object',
309264
definition: 'Returns an object',
310-
}
311-
);
312-
313-
const arrayResult = await app.resolve(createEvent('array-tool'), context);
314-
expect(arrayResult.response.function).toEqual('array-tool');
315-
expect(
316-
arrayResult.response.functionResponse.responseBody.TEXT.body
317-
).toEqual('[1,"two",false,null]');
318-
319-
const objectResult = await app.resolve(createEvent('object-tool'), context);
320-
expect(objectResult.response.function).toEqual('object-tool');
321-
expect(
322-
objectResult.response.functionResponse.responseBody.TEXT.body
323-
).toEqual(
324-
'{"name":"John Doe","age":30,"isActive":true,"address":{"street":"123 Main St","city":"Anytown"}}'
325-
);
326-
});
327-
328-
it('handles functions that return null or undefined by returning a string', async () => {
329-
// Prepare
330-
const app = new BedrockAgentFunctionResolver();
331-
332-
app.tool(async () => null, {
333-
name: 'null-tool',
334-
definition: 'Returns null',
335-
});
336-
337-
app.tool(async () => undefined, {
338-
name: 'undefined-tool',
339-
definition: 'Returns undefined',
340-
});
341-
342-
app.tool(async () => {}, {
343-
name: 'no-return-tool',
344-
definition: 'Has no return statement',
345-
});
346-
347-
const nullResult = await app.resolve(createEvent('null-tool'), context);
348-
expect(nullResult.response.function).toEqual('null-tool');
349-
expect(nullResult.response.functionResponse.responseBody.TEXT.body).toEqual(
350-
''
351-
);
352-
353-
const undefinedResult = await app.resolve(
354-
createEvent('undefined-tool'),
355-
context
356-
);
357-
expect(undefinedResult.response.function).toEqual('undefined-tool');
358-
expect(
359-
undefinedResult.response.functionResponse.responseBody.TEXT.body
360-
).toEqual('');
361-
362-
const noReturnResult = await app.resolve(
363-
createEvent('no-return-tool'),
364-
context
365-
);
366-
expect(noReturnResult.response.function).toEqual('no-return-tool');
367-
expect(
368-
noReturnResult.response.functionResponse.responseBody.TEXT.body
369-
).toEqual('');
370-
});
265+
},
266+
expected:
267+
'{"name":"John Doe","age":30,"isActive":true,"address":{"street":"123 Main St","city":"Anytown"}}',
268+
},
269+
{
270+
toolFunction: async () => [1, 'two', false, null],
271+
toolParams: {
272+
name: 'array',
273+
definition: 'Returns an array',
274+
},
275+
expected: '[1,"two",false,null]',
276+
},
277+
])(
278+
'handles function that returns $toolParams.name',
279+
async ({ toolFunction, toolParams, expected }) => {
280+
// Prepare
281+
const app = new BedrockAgentFunctionResolver();
282+
283+
app.tool(toolFunction, toolParams);
284+
285+
// Act
286+
const actual = await app.resolve(createEvent(toolParams.name), context);
287+
288+
// Asses
289+
expect(actual.response.function).toEqual(toolParams.name);
290+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
291+
expected
292+
);
293+
}
294+
);
295+
296+
it.each([
297+
{
298+
toolFunction: async () => null,
299+
toolParams: {
300+
name: 'null',
301+
definition: 'Returns null',
302+
},
303+
},
304+
{
305+
toolFunction: async () => void 0,
306+
toolParams: {
307+
name: 'undefined',
308+
definition: 'Returns undefined',
309+
},
310+
},
311+
])(
312+
'handles functions that return $toolParams.name by returning an empty string',
313+
async ({ toolFunction, toolParams }) => {
314+
// Prepare
315+
const app = new BedrockAgentFunctionResolver();
316+
317+
app.tool(toolFunction, toolParams);
318+
319+
// Assess
320+
const actual = await app.resolve(createEvent(toolParams.name), context);
321+
322+
// Act
323+
expect(actual.response.function).toEqual(toolParams.name);
324+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
325+
''
326+
);
327+
}
328+
);
371329

372330
it('handles functions that throw errors', async () => {
373331
// Prepare
@@ -383,14 +341,14 @@ describe('Class: BedrockAgentFunctionResolver', () => {
383341
}
384342
);
385343

386-
const errorResult = await app.resolve(
387-
createEvent('error-tool', []),
388-
context
344+
// Act
345+
const actual = await app.resolve(createEvent('error-tool', []), context);
346+
347+
// Assess
348+
expect(actual.response.function).toEqual('error-tool');
349+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
350+
'Error when invoking tool: Error: Something went wrong'
389351
);
390-
expect(errorResult.response.function).toEqual('error-tool');
391-
expect(
392-
errorResult.response.functionResponse.responseBody.TEXT.body
393-
).toEqual('Error when invoking tool: Error: Something went wrong');
394352
expect(console.error).toHaveBeenCalledWith(
395353
'An error occurred in tool error-tool.',
396354
new Error('Something went wrong')
@@ -401,7 +359,6 @@ describe('Class: BedrockAgentFunctionResolver', () => {
401359
// Prepare
402360
const app = new BedrockAgentFunctionResolver();
403361

404-
// Register a tool that returns a simple value
405362
app.tool(
406363
async (params) => {
407364
return `Hello, ${params.name}!`;
@@ -412,7 +369,6 @@ describe('Class: BedrockAgentFunctionResolver', () => {
412369
}
413370
);
414371

415-
// Define custom session attributes and parameters
416372
const customSessionAttrs = {
417373
sessionAttr: '12345',
418374
};
@@ -421,7 +377,6 @@ describe('Class: BedrockAgentFunctionResolver', () => {
421377
promptAttr: 'promptAttr',
422378
};
423379

424-
// Create a custom event with session attributes
425380
const customEvent = {
426381
...createEvent('greeting', [
427382
{
@@ -435,8 +390,10 @@ describe('Class: BedrockAgentFunctionResolver', () => {
435390
promptSessionAttributes: customPromptAttrs,
436391
};
437392

393+
// Act
438394
const result = await app.resolve(customEvent, context);
439395

396+
// Assess
440397
expect(result).toEqual({
441398
messageVersion: '1.0',
442399
response: {

0 commit comments

Comments
 (0)