test: refactor unit tests to shared helpers

Author: marslavish

packages/agent/__tests__/agent.test.ts

@@ -4,117 +4,35 @@ import {
   createAssistantMessageEventStream,
   type ModelDescriptor,
 } from 'agentic-kit';
+import {
+  createScriptedProvider,
+  makeFakeAssistantMessage,
+  makeFakeModel,
+} from '@test/index';
 
 import { Agent } from '../src';
 
-function createModel(): ModelDescriptor {
-  return {
-    id: 'demo',
-    name: 'Demo',
-    api: 'fake',
-    provider: 'fake',
-    baseUrl: 'http://fake.local',
-    input: ['text'],
-    reasoning: false,
-    tools: true,
-  };
-}
-
 describe('@agentic-kit/agent', () => {
   it('runs a minimal sequential tool loop', async () => {
     const responses = [
-      {
-        role: 'assistant' as const,
-        api: 'fake',
-        provider: 'fake',
-        model: 'demo',
-        usage: {
-          input: 1,
-          output: 1,
-          cacheRead: 0,
-          cacheWrite: 0,
-          totalTokens: 2,
-          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-        },
-        stopReason: 'toolUse' as const,
-        timestamp: Date.now(),
+      makeFakeAssistantMessage({
+        usage: makeUsage(),
+        stopReason: 'toolUse',
         content: [
-          { type: 'toolCall' as const, id: 'tool_1', name: 'echo', arguments: { text: 'hello' } },
+          { type: 'toolCall', id: 'tool_1', name: 'echo', arguments: { text: 'hello' } },
         ],
-      },
-      {
-        role: 'assistant' as const,
-        api: 'fake',
-        provider: 'fake',
-        model: 'demo',
-        usage: {
-          input: 1,
-          output: 1,
-          cacheRead: 0,
-          cacheWrite: 0,
-          totalTokens: 2,
-          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-        },
-        stopReason: 'stop' as const,
-        timestamp: Date.now(),
-        content: [{ type: 'text' as const, text: 'done' }],
-      },
+      }),
+      makeFakeAssistantMessage({
+        usage: makeUsage(),
+        stopReason: 'stop',
+        content: [{ type: 'text', text: 'done' }],
+      }),
     ];
 
-    let callIndex = 0;
-    const streamFn = (_model: ModelDescriptor, _context: Context) => {
-      const stream = createAssistantMessageEventStream();
-      const response = responses[callIndex++];
-
-      queueMicrotask(() => {
-        stream.push({ type: 'start', partial: response });
-        if (response.content[0].type === 'toolCall') {
-          stream.push({
-            type: 'toolcall_start',
-            contentIndex: 0,
-            partial: response,
-          });
-          stream.push({
-            type: 'toolcall_end',
-            contentIndex: 0,
-            toolCall: response.content[0],
-            partial: response,
-          });
-        } else {
-          stream.push({
-            type: 'text_start',
-            contentIndex: 0,
-            partial: response,
-          });
-          stream.push({
-            type: 'text_delta',
-            contentIndex: 0,
-            delta: 'done',
-            partial: response,
-          });
-          stream.push({
-            type: 'text_end',
-            contentIndex: 0,
-            content: 'done',
-            partial: response,
-          });
-        }
-        stream.push({
-          type: 'done',
-          reason: response.stopReason === 'toolUse' ? 'toolUse' : 'stop',
-          message: response,
-        });
-        stream.end(response);
-      });
-
-      return stream;
-    };
-
+    const provider = createScriptedProvider({ responses });
     const agent = new Agent({
-      initialState: {
-        model: createModel(),
-      },
-      streamFn,
+      initialState: { model: makeFakeModel({ id: 'demo', name: 'Demo' }) },
+      streamFn: provider.stream,
     });
 
     agent.setTools([
@@ -155,20 +73,20 @@ describe('@agentic-kit/agent', () => {
 
   it('turns tool argument validation failures into error tool results and continues', async () => {
     const responses = [
-      createAssistantResponse({
+      makeFakeAssistantMessage({
         stopReason: 'toolUse',
         content: [{ type: 'toolCall', id: 'tool_1', name: 'echo', arguments: {} }],
       }),
-      createAssistantResponse({
+      makeFakeAssistantMessage({
         stopReason: 'stop',
         content: [{ type: 'text', text: 'recovered' }],
       }),
     ];
 
-    let callIndex = 0;
+    const provider = createScriptedProvider({ responses });
     const agent = new Agent({
-      initialState: { model: createModel() },
-      streamFn: () => streamMessage(responses[callIndex++]),
+      initialState: { model: makeFakeModel({ id: 'demo', name: 'Demo' }) },
+      streamFn: provider.stream,
     });
 
     const execute = jest.fn(async () => ({
@@ -211,10 +129,10 @@ describe('@agentic-kit/agent', () => {
 
   it('records aborted assistant turns when the active stream is cancelled', async () => {
     const agent = new Agent({
-      initialState: { model: createModel() },
+      initialState: { model: makeFakeModel({ id: 'demo', name: 'Demo' }) },
       streamFn: (_model: ModelDescriptor, _context: Context, options) => {
         const stream = createAssistantMessageEventStream();
-        const partial = createAssistantResponse({
+        const partial = makeFakeAssistantMessage({
           stopReason: 'stop',
           content: [{ type: 'text', text: '' }],
         });
@@ -225,7 +143,7 @@ describe('@agentic-kit/agent', () => {
           options?.signal?.addEventListener(
             'abort',
             () => {
-              const aborted = createAssistantResponse({
+              const aborted: AssistantMessage = makeFakeAssistantMessage({
                 stopReason: 'aborted',
                 errorMessage: 'aborted by test',
                 content: [],
@@ -256,76 +174,13 @@ describe('@agentic-kit/agent', () => {
   });
 });
 
-function createAssistantResponse(overrides: Partial<AssistantMessage>): AssistantMessage {
+function makeUsage() {
   return {
-    ...createAssistantResponseBase(),
-    ...overrides,
+    input: 1,
+    output: 1,
+    cacheRead: 0,
+    cacheWrite: 0,
+    totalTokens: 2,
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
   };
 }
-
-function createAssistantResponseBase(): AssistantMessage {
-  return {
-    role: 'assistant' as const,
-    api: 'fake',
-    provider: 'fake',
-    model: 'demo',
-    usage: {
-      input: 1,
-      output: 1,
-      cacheRead: 0,
-      cacheWrite: 0,
-      totalTokens: 2,
-      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-    },
-    stopReason: 'stop' as const,
-    timestamp: Date.now(),
-    content: [] as AssistantMessage['content'],
-  };
-}
-
-function streamMessage(message: AssistantMessage) {
-  const stream = createAssistantMessageEventStream();
-
-  queueMicrotask(() => {
-    stream.push({ type: 'start', partial: message });
-    if (message.content[0]?.type === 'toolCall') {
-      stream.push({
-        type: 'toolcall_start',
-        contentIndex: 0,
-        partial: message,
-      });
-      stream.push({
-        type: 'toolcall_end',
-        contentIndex: 0,
-        toolCall: message.content[0],
-        partial: message,
-      });
-    } else {
-      stream.push({
-        type: 'text_start',
-        contentIndex: 0,
-        partial: message,
-      });
-      stream.push({
-        type: 'text_delta',
-        contentIndex: 0,
-        delta: message.content[0]?.type === 'text' ? message.content[0].text : '',
-        partial: message,
-      });
-      stream.push({
-        type: 'text_end',
-        contentIndex: 0,
-        content: message.content[0]?.type === 'text' ? message.content[0].text : '',
-        partial: message,
-      });
-    }
-    stream.push({
-      type: 'done',
-      reason: message.stopReason === 'toolUse' ? 'toolUse' : 'stop',
-      message,
-    });
-    stream.end(message);
-  });
-
-  return stream;
-}