From b5f26df835c3a08e9065f327afc4a6bc1997c9ab Mon Sep 17 00:00:00 2001
From: Mike Donnalley <mdonnalley@salesforce.com>
Date: Wed, 4 Dec 2024 13:50:53 -0700
Subject: [PATCH] feat: export test formatters

---
 src/agentTester.ts       | 30 +++++++++---------------------
 src/index.ts             |  2 ++
 test/agentTester.test.ts | 20 ++++----------------
 3 files changed, 15 insertions(+), 37 deletions(-)

diff --git a/src/agentTester.ts b/src/agentTester.ts
index 86198f6..5872e4f 100644
--- a/src/agentTester.ts
+++ b/src/agentTester.ts
@@ -8,8 +8,6 @@ import { Connection, Lifecycle, PollingClient, StatusResult } from '@salesforce/
 import { Duration } from '@salesforce/kit';
 import { MaybeMock } from './maybe-mock';
 
-type Format = 'human' | 'json';
-
 export type TestStatus = 'NEW' | 'IN_PROGRESS' | 'COMPLETED' | 'ERROR';
 
 export type AgentTestStartResponse = {
@@ -91,25 +89,22 @@ export class AgentTester {
   public async poll(
     jobId: string,
     {
-      format = 'human',
       timeout = Duration.minutes(5),
     }: {
-      format?: Format;
       timeout?: Duration;
     } = {
-      format: 'human',
       timeout: Duration.minutes(5),
     }
-  ): Promise<{ response: AgentTestDetailsResponse; formatted: string }> {
+  ): Promise<AgentTestDetailsResponse> {
     const lifecycle = Lifecycle.getInstance();
     const client = await PollingClient.create({
       poll: async (): Promise<StatusResult> => {
         // NOTE: we don't actually need to call the status API here since all the same information is present on the
         // details API. We could just call the details API and check the status there.
-        const [detailsResponse, statusResponse] = await Promise.all([this.details(jobId, format), this.status(jobId)]);
-        const totalTestCases = detailsResponse.response.testCases.length;
-        const failingTestCases = detailsResponse.response.testCases.filter((tc) => tc.status === 'ERROR').length;
-        const passingTestCases = detailsResponse.response.testCases.filter(
+        const [detailsResponse, statusResponse] = await Promise.all([this.details(jobId), this.status(jobId)]);
+        const totalTestCases = detailsResponse.testCases.length;
+        const failingTestCases = detailsResponse.testCases.filter((tc) => tc.status === 'ERROR').length;
+        const passingTestCases = detailsResponse.testCases.filter(
           (tc) => tc.status === 'COMPLETED' && tc.expectationResults.every((r) => r.result === 'Passed')
         ).length;
 
@@ -121,7 +116,7 @@ export class AgentTester {
             failingTestCases,
             passingTestCases,
           });
-          return { payload: await this.details(jobId, format), completed: true };
+          return { payload: detailsResponse, completed: true };
         }
 
         await lifecycle.emit('AGENT_TEST_POLLING_EVENT', {
@@ -137,21 +132,14 @@ export class AgentTester {
       timeout,
     });
 
-    const result = await client.subscribe<{ response: AgentTestDetailsResponse; formatted: string }>();
+    const result = await client.subscribe<AgentTestDetailsResponse>();
     return result;
   }
 
-  public async details(
-    jobId: string,
-    format: Format = 'human'
-  ): Promise<{ response: AgentTestDetailsResponse; formatted: string }> {
+  public async details(jobId: string): Promise<AgentTestDetailsResponse> {
     const url = `/einstein/ai-evaluations/runs/${jobId}/details`;
 
-    const response = await this.maybeMock.request<AgentTestDetailsResponse>('GET', url);
-    return {
-      response,
-      formatted: format === 'human' ? await humanFormat(jobId, response) : await jsonFormat(response),
-    };
+    return this.maybeMock.request<AgentTestDetailsResponse>('GET', url);
   }
 
   public async cancel(jobId: string): Promise<{ success: boolean }> {
diff --git a/src/index.ts b/src/index.ts
index 51caddd..3be5846 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -16,6 +16,8 @@ export {
 export { Agent } from './agent';
 export {
   AgentTester,
+  humanFormat,
+  jsonFormat,
   type AgentTestDetailsResponse,
   type AgentTestStartResponse,
   type AgentTestStatusResponse,
diff --git a/test/agentTester.test.ts b/test/agentTester.test.ts
index 9b1526f..704b66a 100644
--- a/test/agentTester.test.ts
+++ b/test/agentTester.test.ts
@@ -51,25 +51,13 @@ describe('AgentTester', () => {
   });
 
   describe('poll', () => {
-    it('should poll until test run is complete (human format)', async () => {
+    it('should poll until test run is complete', async () => {
       const tester = new AgentTester(connection);
       await tester.start('suiteId');
-      const output = await tester.poll('4KBSM000000003F4AQ');
-      expect(output).to.be.ok;
-      // TODO: make these assertions more meaningful
-      expect(output.formatted).to.include('Test Case #1');
-      expect(output.formatted).to.include('Test Case #2');
-      expect(output.response.testCases[0].status).to.equal('COMPLETED');
-    });
-
-    it('should poll until test run is complete (json format)', async () => {
-      const tester = new AgentTester(connection);
-      await tester.start('suiteId');
-      const output = await tester.poll('4KBSM000000003F4AQ', { format: 'json' });
-      expect(output).to.be.ok;
+      const response = await tester.poll('4KBSM000000003F4AQ');
+      expect(response).to.be.ok;
       // TODO: make these assertions more meaningful
-      expect(JSON.parse(output.formatted)).to.deep.equal(output.response);
-      expect(output.response.testCases[0].status).to.equal('COMPLETED');
+      expect(response.testCases[0].status).to.equal('COMPLETED');
     });
   });