From 47d7b31fe4c6ebfd6f4873f8b7698ec10b569e4b Mon Sep 17 00:00:00 2001 From: Varun Nuthalapati Date: Sun, 28 Jun 2026 12:39:34 -0700 Subject: [PATCH] test: add unit tests for TypeScript ChainAgent and BedrockFlowsAgent Adds ChainAgent.test.ts (12 tests) covering constructor validation, the chaining pipeline, passthrough of additionalParams, streaming on the last agent, error propagation, and defaultOutput fallback. Adds BedrockFlowsAgent.test.ts (8 tests) covering constructor options, processRequest happy path, custom encoder/decoder, missing response stream, and client error wrapping. Closes #539 --- .../tests/agents/BedrockFlowsAgent.test.ts | 119 +++++++++++++++ typescript/tests/agents/ChainAgent.test.ts | 136 ++++++++++++++++++ 2 files changed, 255 insertions(+) create mode 100644 typescript/tests/agents/BedrockFlowsAgent.test.ts create mode 100644 typescript/tests/agents/ChainAgent.test.ts diff --git a/typescript/tests/agents/BedrockFlowsAgent.test.ts b/typescript/tests/agents/BedrockFlowsAgent.test.ts new file mode 100644 index 000000000..be777128e --- /dev/null +++ b/typescript/tests/agents/BedrockFlowsAgent.test.ts @@ -0,0 +1,119 @@ +import { BedrockFlowsAgent, BedrockFlowsAgentOptions } from '../../src/agents/bedrockFlowsAgent'; +import { BedrockAgentRuntimeClient, InvokeFlowCommand } from '@aws-sdk/client-bedrock-agent-runtime'; +import { ConversationMessage, ParticipantRole } from '../../src/types'; + +jest.mock('@aws-sdk/client-bedrock-agent-runtime'); +jest.mock('../../src/common/src/awsSdkUtils', () => ({ + addUserAgentMiddleware: jest.fn(), +})); + +describe('BedrockFlowsAgent', () => { + let mockClient: jest.Mocked; + + const defaultOptions: BedrockFlowsAgentOptions = { + name: 'TestFlowsAgent', + description: 'Test BedrockFlowsAgent', + flowIdentifier: 'flow-123', + flowAliasIdentifier: 'alias-456', + }; + + function makeResponseStream(document: string) { + async function* gen() { + yield { flowOutputEvent: { content: { document } } }; + } + return { responseStream: gen() }; + } + + beforeEach(() => { + jest.clearAllMocks(); + mockClient = { send: jest.fn(), use: jest.fn(), destroy: jest.fn() } as any; + (BedrockAgentRuntimeClient as jest.Mock).mockImplementation(() => mockClient); + (InvokeFlowCommand as unknown as jest.Mock).mockImplementation((p) => p); + }); + + describe('constructor', () => { + it('creates a BedrockAgentRuntimeClient with no arguments when no region given', () => { + new BedrockFlowsAgent(defaultOptions); + expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith(); + }); + + it('creates a BedrockAgentRuntimeClient with specified region', () => { + new BedrockFlowsAgent({ ...defaultOptions, region: 'eu-west-1' }); + expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({ region: 'eu-west-1' }); + }); + + it('uses a pre-built client when bedrockAgentClient is provided', () => { + const prebuilt = { send: jest.fn() } as any; + const agent = new BedrockFlowsAgent({ ...defaultOptions, bedrockAgentClient: prebuilt }); + expect(agent).toBeDefined(); + expect(BedrockAgentRuntimeClient).not.toHaveBeenCalled(); + }); + + it('defaults enableTrace to false', () => { + const agent = new BedrockFlowsAgent(defaultOptions); + expect((agent as any).enableTrace).toBe(false); + }); + + it('sets enableTrace when provided', () => { + const agent = new BedrockFlowsAgent({ ...defaultOptions, enableTrace: true }); + expect((agent as any).enableTrace).toBe(true); + }); + }); + + describe('processRequest', () => { + const userId = 'user1'; + const sessionId = 'session1'; + const chatHistory: ConversationMessage[] = []; + + it('invokes the flow and returns the decoded response', async () => { + mockClient.send.mockResolvedValueOnce(makeResponseStream('Flow answer') as never); + + const agent = new BedrockFlowsAgent(defaultOptions); + const result = await agent.processRequest('What is 2+2?', userId, sessionId, chatHistory); + + expect(mockClient.send).toHaveBeenCalledTimes(1); + expect(result).toEqual({ + role: ParticipantRole.ASSISTANT, + content: [{ text: 'Flow answer' }], + }); + }); + + it('throws when the response stream is missing', async () => { + mockClient.send.mockResolvedValueOnce({ responseStream: null } as never); + + const agent = new BedrockFlowsAgent(defaultOptions); + await expect(agent.processRequest('input', userId, sessionId, chatHistory)) + .rejects.toThrow(/Error processing request with Bedrock/); + }); + + it('uses a custom flowInputEncoder when provided', async () => { + const encoder = jest.fn().mockReturnValue('encoded input'); + mockClient.send.mockResolvedValueOnce(makeResponseStream('ok') as never); + + const agent = new BedrockFlowsAgent({ ...defaultOptions, flowInputEncoder: encoder }); + await agent.processRequest('hello', userId, sessionId, chatHistory); + + expect(encoder).toHaveBeenCalledWith(agent, 'hello', expect.objectContaining({ userId, sessionId })); + }); + + it('uses a custom flowOutputDecoder when provided', async () => { + const decoded: ConversationMessage = { role: ParticipantRole.ASSISTANT, content: [{ text: 'custom' }] }; + const decoder = jest.fn().mockReturnValue(decoded); + mockClient.send.mockResolvedValueOnce(makeResponseStream('raw') as never); + + const agent = new BedrockFlowsAgent({ ...defaultOptions, flowOutputDecoder: decoder }); + const result = await agent.processRequest('input', userId, sessionId, chatHistory); + + expect(decoder).toHaveBeenCalledWith(agent, 'raw'); + expect(result).toEqual(decoded); + }); + + it('wraps client errors in a descriptive error message', async () => { + mockClient.send.mockRejectedValueOnce(new Error('Network failure') as never); + + const agent = new BedrockFlowsAgent(defaultOptions); + await expect(agent.processRequest('input', userId, sessionId, chatHistory)) + .rejects.toThrow('Error processing request with Bedrock: Network failure'); + }); + }); +}); diff --git a/typescript/tests/agents/ChainAgent.test.ts b/typescript/tests/agents/ChainAgent.test.ts new file mode 100644 index 000000000..d0ac57e82 --- /dev/null +++ b/typescript/tests/agents/ChainAgent.test.ts @@ -0,0 +1,136 @@ +import { ChainAgent, ChainAgentOptions } from '../../src/agents/chainAgent'; +import { Agent, AgentOptions } from '../../src/agents/agent'; +import { ConversationMessage, ParticipantRole } from '../../src/types'; + +// Minimal concrete Agent subclass for testing +class MockAgent extends Agent { + private response: ConversationMessage | AsyncIterable; + + constructor( + options: AgentOptions, + response: ConversationMessage | AsyncIterable + ) { + super(options); + this.response = response; + } + + async processRequest(): Promise> { + return this.response; + } +} + +function makeMessage(text: string): ConversationMessage { + return { role: ParticipantRole.ASSISTANT, content: [{ text }] }; +} + +function makeAgent(name: string, response: ConversationMessage | AsyncIterable): MockAgent { + return new MockAgent({ name, description: `${name} description` }, response); +} + +describe('ChainAgent', () => { + const defaultOptions: ChainAgentOptions = { + name: 'TestChain', + description: 'Test chain agent', + agents: [makeAgent('Agent1', makeMessage('step one output'))], + }; + + describe('constructor', () => { + it('throws when agents array is empty', () => { + expect(() => new ChainAgent({ ...defaultOptions, agents: [] })).toThrow( + 'ChainAgent requires at least one agent in the chain.' + ); + }); + + it('initialises with a single agent', () => { + const chain = new ChainAgent(defaultOptions); + expect(chain.agents).toHaveLength(1); + }); + + it('uses custom defaultOutput when provided', () => { + const chain = new ChainAgent({ ...defaultOptions, defaultOutput: 'custom fallback' }); + expect(chain).toBeDefined(); + }); + }); + + describe('processRequest', () => { + const userId = 'user1'; + const sessionId = 'session1'; + const chatHistory: ConversationMessage[] = []; + + it('returns the output of a single-agent chain', async () => { + const chain = new ChainAgent(defaultOptions); + const result = await chain.processRequest('hello', userId, sessionId, chatHistory) as ConversationMessage; + + expect(result.role).toBe(ParticipantRole.ASSISTANT); + expect(result.content[0]).toHaveProperty('text', 'step one output'); + }); + + it('pipes output of each agent as input to the next', async () => { + const spy1 = jest.fn().mockResolvedValue(makeMessage('output of agent 1')); + const spy2 = jest.fn().mockResolvedValue(makeMessage('output of agent 2')); + + const agent1 = new MockAgent({ name: 'A1', description: 'd' }, makeMessage('')); + const agent2 = new MockAgent({ name: 'A2', description: 'd' }, makeMessage('')); + agent1.processRequest = spy1 as any; + agent2.processRequest = spy2 as any; + + const chain = new ChainAgent({ ...defaultOptions, agents: [agent1, agent2] }); + await chain.processRequest('initial input', userId, sessionId, chatHistory); + + expect(spy1).toHaveBeenCalledWith('initial input', userId, sessionId, chatHistory, undefined); + expect(spy2).toHaveBeenCalledWith('output of agent 1', userId, sessionId, chatHistory, undefined); + }); + + it('passes additionalParams through to each agent', async () => { + const spy = jest.fn().mockResolvedValue(makeMessage('ok')); + const agent = new MockAgent({ name: 'A', description: 'd' }, makeMessage('')); + agent.processRequest = spy as any; + + const chain = new ChainAgent({ ...defaultOptions, agents: [agent] }); + const params = { key: 'val' }; + await chain.processRequest('input', userId, sessionId, chatHistory, params); + + expect(spy).toHaveBeenCalledWith('input', userId, sessionId, chatHistory, params); + }); + + it('returns default response when an intermediate agent returns no text content', async () => { + const emptyAgent = makeAgent('Empty', { role: ParticipantRole.ASSISTANT, content: [] }); + const secondAgent = makeAgent('Second', makeMessage('should not reach')); + + const chain = new ChainAgent({ ...defaultOptions, agents: [emptyAgent, secondAgent] }); + const result = await chain.processRequest('input', userId, sessionId, chatHistory) as ConversationMessage; + + expect(result.content[0]).toHaveProperty('text', 'No output generated from the chain.'); + }); + + it('throws when an agent in the chain errors', async () => { + const errorAgent = new MockAgent({ name: 'ErrorAgent', description: 'd' }, makeMessage('')); + errorAgent.processRequest = jest.fn().mockRejectedValue(new Error('downstream failure')) as any; + + const chain = new ChainAgent({ ...defaultOptions, agents: [errorAgent] }); + await expect(chain.processRequest('input', userId, sessionId, chatHistory)) + .rejects.toMatch(/Error processing request with agent ErrorAgent/); + }); + + it('allows the last agent to return a streaming response', async () => { + async function* streamGen() { yield 'chunk'; } + const streamingAgent = makeAgent('Streamer', streamGen()); + + const chain = new ChainAgent({ ...defaultOptions, agents: [streamingAgent] }); + const result = await chain.processRequest('input', userId, sessionId, chatHistory); + + expect(typeof (result as any)[Symbol.asyncIterator]).toBe('function'); + }); + + it('returns default response when a non-last agent returns a streaming response', async () => { + async function* streamGen() { yield 'chunk'; } + const streamingAgent = makeAgent('Streamer', streamGen()); + const nextAgent = makeAgent('Next', makeMessage('next')); + + const chain = new ChainAgent({ ...defaultOptions, agents: [streamingAgent, nextAgent] }); + const result = await chain.processRequest('input', userId, sessionId, chatHistory) as ConversationMessage; + + expect(result.content[0]).toHaveProperty('text', 'No output generated from the chain.'); + }); + }); +});