[AI Assistance] Add unit test helpers

Bug: none
Change-Id: I1a81e80413d19795b4582a5780c09201d836f802
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6236124
Commit-Queue: Alex Rudenko <[email protected]>
Reviewed-by: Nikolay Vitkov <[email protected]>

Author: OrKoN

front_end/panels/ai_assistance/agents/AiAgent.test.ts

@@ -3,23 +3,14 @@
 // found in the LICENSE file.
 
 import * as Host from '../../../core/host/host.js';
+import {mockAidaClient} from '../../../testing/AiAssistanceHelpers.js';
 import {
   describeWithEnvironment,
 } from '../../../testing/EnvironmentHelpers.js';
 import * as AiAssistance from '../ai_assistance.js';
 
 const {AiAgent, ResponseType, ConversationContext, ErrorType} = AiAssistance;
 
-function mockAidaClient(
-    fetch: (_: Host.AidaClient.AidaRequest, options?: {signal: AbortSignal}) =>
-        AsyncGenerator<Host.AidaClient.AidaResponse, void, void>,
-    ): Host.AidaClient.AidaClient {
-  return {
-    fetch,
-    registerClientEvent: () => Promise.resolve({}),
-  };
-}
-
 function mockConversationContext(): AiAssistance.ConversationContext<unknown> {
   return new (class extends ConversationContext<unknown>{
     override getOrigin(): string {
@@ -70,7 +61,7 @@ describeWithEnvironment('AiAgent', () => {
 
     it('builds a request with a temperature', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
       });
       assert.strictEqual(
           agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER).options?.temperature,
@@ -80,15 +71,15 @@ describeWithEnvironment('AiAgent', () => {
 
     it('builds a request with a temperature -1', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
       });
       agent.options.temperature = -1;
       assert.isUndefined(agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER).options?.temperature);
     });
 
     it('builds a request with a model id', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
       });
       assert.strictEqual(
           agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER).options?.model_id,
@@ -98,7 +89,7 @@ describeWithEnvironment('AiAgent', () => {
 
     it('builds a request with logging', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
         serverSideLoggingEnabled: true,
       });
       assert.isFalse(
@@ -107,7 +98,7 @@ describeWithEnvironment('AiAgent', () => {
 
     it('builds a request without logging', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
         serverSideLoggingEnabled: false,
       });
       assert.isTrue(agent
@@ -121,7 +112,7 @@ describeWithEnvironment('AiAgent', () => {
 
     it('builds a request with input', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
         serverSideLoggingEnabled: false,
       });
       const request = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
@@ -131,15 +122,15 @@ describeWithEnvironment('AiAgent', () => {
 
     it('builds a request with a sessionId', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
       });
       const request = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
       assert.strictEqual(request.metadata?.string_session_id, 'sessionId');
     });
 
     it('builds a request with preamble', async () => {
       const agent = new AiAgentMock({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
       });
       const request = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
       assert.deepEqual(request.current_message?.parts[0], {text: 'test input'});
@@ -164,7 +155,7 @@ describeWithEnvironment('AiAgent', () => {
       }
 
       const agent = new AiAgentMockWithoutPreamble({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
       });
       const request = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
       assert.deepEqual(request.current_message?.parts[0], {text: 'test input'});
@@ -173,21 +164,10 @@ describeWithEnvironment('AiAgent', () => {
     });
 
     it('builds a request with chat history', async () => {
-      let count = 0;
-      async function* generateAndAnswer() {
-        if (count === 0) {
-          yield {
-            explanation: 'answer',
-            metadata: {},
-            completed: true,
-          };
-        }
-
-        count++;
-      }
-
       const agent = new AiAgentMock({
-        aidaClient: mockAidaClient(generateAndAnswer),
+        aidaClient: mockAidaClient([[{
+          explanation: 'answer',
+        }]]),
         serverSideLoggingEnabled: true,
       });
       await Array.fromAsync(agent.run('question', {selected: null}));
@@ -207,26 +187,10 @@ describeWithEnvironment('AiAgent', () => {
     });
 
     it('builds a request with aborted query in history', async () => {
-      let count = 0;
-      async function* generateAndAnswer(_: Host.AidaClient.AidaRequest, options?: {signal?: AbortSignal}) {
-        if (options?.signal?.aborted) {
-          // Should this be the aborted Error?
-          throw new Error('Aborted');
-        }
-
-        if (count === 0) {
-          yield {
-            explanation: 'answer',
-            metadata: {},
-            completed: true,
-          };
-        }
-
-        count++;
-      }
-
       const agent = new AiAgentMock({
-        aidaClient: mockAidaClient(generateAndAnswer),
+        aidaClient: mockAidaClient([[{
+          explanation: 'answer',
+        }]]),
         serverSideLoggingEnabled: true,
       });
 
@@ -243,21 +207,15 @@ describeWithEnvironment('AiAgent', () => {
   describe('run', () => {
     describe('partial yielding for answers', () => {
       it('should yield partial answer with final answer at the end', async () => {
-        async function* generateAnswerAfterPartial() {
-          yield {
-            explanation: 'Partial ans',
-            metadata: {},
-            completed: false,
-          };
-
-          yield {
-            explanation: 'Partial answer is now completed',
-            metadata: {},
-            completed: true,
-          };
-        }
         const agent = new AiAgentMock({
-          aidaClient: mockAidaClient(generateAnswerAfterPartial),
+          aidaClient: mockAidaClient([[
+            {
+              explanation: 'Partial ans',
+            },
+            {
+              explanation: 'Partial answer is now completed',
+            }
+          ]]),
         });
 
         const responses = await Array.fromAsync(agent.run('query', {selected: mockConversationContext()}));
@@ -284,21 +242,15 @@ describeWithEnvironment('AiAgent', () => {
       });
 
       it('should not add partial answers to history', async () => {
-        async function* generateAnswerAfterPartial() {
-          yield {
-            explanation: 'Partial ans',
-            metadata: {},
-            completed: false,
-          };
-
-          yield {
-            explanation: 'Partial answer is now completed',
-            metadata: {},
-            completed: true,
-          };
-        }
         const agent = new AiAgentMock({
-          aidaClient: mockAidaClient(generateAnswerAfterPartial),
+          aidaClient: mockAidaClient([[
+            {
+              explanation: 'Partial ans',
+            },
+            {
+              explanation: 'Partial answer is now completed',
+            }
+          ]]),
         });
 
         await Array.fromAsync(agent.run('query', {selected: mockConversationContext()}));
@@ -317,10 +269,8 @@ describeWithEnvironment('AiAgent', () => {
     });
 
     it('should yield unknown error when aidaFetch does not return anything', async () => {
-      async function* generateNothing() {
-      }
       const agent = new AiAgentMock({
-        aidaClient: mockAidaClient(generateNothing),
+        aidaClient: mockAidaClient([]),
       });
 
       const responses = await Array.fromAsync(agent.run('query', {selected: mockConversationContext()}));
@@ -343,27 +293,15 @@ describeWithEnvironment('AiAgent', () => {
 
   describe('runFromHistory', () => {
     it('should run', async () => {
-      let count = 0;
-      async function* generateAndAnswer() {
-        if (count === 0) {
-          yield {
+      const agent = new AiAgentMock({
+        aidaClient: mockAidaClient([
+          [{
             explanation: 'first answer',
-            metadata: {},
-            completed: true,
-          };
-        } else {
-          yield {
+          }],
+          [{
             explanation: 'second answer',
-            metadata: {},
-            completed: true,
-          };
-        }
-
-        count++;
-      }
-
-      const agent = new AiAgentMock({
-        aidaClient: mockAidaClient(generateAndAnswer),
+          }]
+        ]),
         serverSideLoggingEnabled: true,
       });
       await Array.fromAsync(agent.run('first question', {selected: null}));
@@ -495,7 +433,7 @@ describeWithEnvironment('AiAgent', () => {
 
     it('should build a request with functions', () => {
       const agent = new AgentWithFunction({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient(),
       });
       agent.options.temperature = -1;
       assert.deepEqual(

front_end/panels/ai_assistance/agents/FileAgent.test.ts

@@ -10,6 +10,7 @@ import type * as Protocol from '../../../generated/protocol.js';
 import * as Bindings from '../../../models/bindings/bindings.js';
 import * as Logs from '../../../models/logs/logs.js';
 import * as Workspace from '../../../models/workspace/workspace.js';
+import {mockAidaClient} from '../../../testing/AiAssistanceHelpers.js';
 import {
   createTarget,
   getGetHostConfigStub,
@@ -31,15 +32,6 @@ describeWithMockConnection('FileAgent', () => {
     });
   }
 
-  function mockAidaClient(
-      fetch: () => AsyncGenerator<Host.AidaClient.AidaResponse, void, void>,
-      ): Host.AidaClient.AidaClient {
-    return {
-      fetch,
-      registerClientEvent: () => Promise.resolve({}),
-    };
-  }
-
   beforeEach(() => {
     const workspace = Workspace.Workspace.WorkspaceImpl.instance();
     const targetManager = SDK.TargetManager.TargetManager.instance();
@@ -86,22 +78,8 @@ describeWithMockConnection('FileAgent', () => {
     it('structure matches the snapshot', async () => {
       mockHostConfig('test model');
       sinon.stub(crypto, 'randomUUID').returns('sessionId' as `${string}-${string}-${string}-${string}-${string}`);
-
-      let count = 0;
-      async function* generateAndAnswer() {
-        if (count === 0) {
-          yield {
-            explanation: 'answer',
-            metadata: {},
-            completed: true,
-          };
-        }
-
-        count++;
-      }
-
       const agent = new FileAgent({
-        aidaClient: mockAidaClient(generateAndAnswer),
+        aidaClient: mockAidaClient([[{explanation: 'answer'}]]),
         serverSideLoggingEnabled: true,
       });
       sinon.stub(agent, 'preamble').value('preamble');
@@ -228,18 +206,13 @@ describeWithMockConnection('FileAgent', () => {
 
     testArguments.forEach(args => {
       it('generates an answer ' + args.name, async () => {
-        async function* generateAnswer() {
-          yield {
+        const agent = new FileAgent({
+          aidaClient: mockAidaClient([[{
             explanation: 'This is the answer',
             metadata: {
               rpcGlobalId: 123,
             },
-            completed: true,
-          };
-        }
-
-        const agent = new FileAgent({
-          aidaClient: mockAidaClient(generateAnswer),
+          }]]),
         });
 
         const uiSourceCode = await createUISourceCode({

front_end/panels/ai_assistance/agents/NetworkAgent.test.ts

@@ -7,6 +7,7 @@ import * as Platform from '../../../core/platform/platform.js';
 import * as SDK from '../../../core/sdk/sdk.js';
 import type * as Protocol from '../../../generated/protocol.js';
 import * as Logs from '../../../models/logs/logs.js';
+import {mockAidaClient} from '../../../testing/AiAssistanceHelpers.js';
 import {
   getGetHostConfigStub,
 } from '../../../testing/EnvironmentHelpers.js';
@@ -142,28 +143,14 @@ describeWithMockConnection('NetworkAgent', () => {
       sinon.restore();
     });
 
-    function mockAidaClient(
-        fetch: () => AsyncGenerator<Host.AidaClient.AidaResponse, void, void>,
-        ): Host.AidaClient.AidaClient {
-      return {
-        fetch,
-        registerClientEvent: () => Promise.resolve({}),
-      };
-    }
-
     it('generates an answer', async () => {
-      async function* generateAnswer() {
-        yield {
+      const agent = new NetworkAgent({
+        aidaClient: mockAidaClient([[{
           explanation: 'This is the answer',
           metadata: {
             rpcGlobalId: 123,
           },
-          completed: true,
-        };
-      }
-
-      const agent = new NetworkAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+        }]]),
       });
 
       const responses =