From 983514ee11362c5efe4cdb59802b3ff402b61ef2 Mon Sep 17 00:00:00 2001
From: Joshua Li <joshuali925@gmail.com>
Date: Fri, 7 Jun 2024 19:44:58 +0000
Subject: [PATCH] add unit tests for agents

Signed-off-by: Joshua Li <joshuali925@gmail.com>
---
 .../server/routes/query_assist/agents.test.ts | 102 ++++++++++++++++++
 1 file changed, 102 insertions(+)
 create mode 100644 plugins-extra/query_enhancements/server/routes/query_assist/agents.test.ts

diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/agents.test.ts b/plugins-extra/query_enhancements/server/routes/query_assist/agents.test.ts
new file mode 100644
index 000000000000..a8b5e67153a6
--- /dev/null
+++ b/plugins-extra/query_enhancements/server/routes/query_assist/agents.test.ts
@@ -0,0 +1,102 @@
+import { ApiResponse } from '@opensearch-project/opensearch';
+import { ResponseError } from '@opensearch-project/opensearch/lib/errors';
+import { RequestHandlerContext } from 'src/core/server';
+import { CoreRouteHandlerContext } from '../../../../../src/core/server/core_route_handler_context';
+import { loggerMock } from '../../../../../src/core/server/logging/logger.mock';
+import { coreMock, httpServerMock } from '../../../../../src/core/server/mocks';
+import { getAgentIdByConfig, requestAgentByConfig } from './agents';
+
+describe('Agents helper functions', () => {
+  const coreContext = new CoreRouteHandlerContext(
+    coreMock.createInternalStart(),
+    httpServerMock.createOpenSearchDashboardsRequest()
+  );
+  const client = coreContext.opensearch.client.asCurrentUser;
+  const mockedTransport = client.transport.request as jest.Mock;
+  const context: RequestHandlerContext = {
+    core: coreContext,
+    dataSource: jest.fn(),
+    query_assist: { dataSourceEnabled: false, logger: loggerMock.create() },
+  };
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('searches agent id by name', async () => {
+    mockedTransport.mockResolvedValueOnce({
+      body: {
+        type: 'agent',
+        configuration: { agent_id: 'agentId' },
+      },
+    });
+    const id = await getAgentIdByConfig(client, 'test_agent');
+    expect(id).toEqual('agentId');
+    expect(mockedTransport.mock.calls[0]).toMatchInlineSnapshot(`
+      Array [
+        Object {
+          "method": "GET",
+          "path": "/_plugins/_ml/config/test_agent",
+        },
+      ]
+    `);
+  });
+
+  it('handles not found errors', async () => {
+    mockedTransport.mockRejectedValueOnce(
+      new ResponseError(({
+        body: {
+          error: {
+            root_cause: [
+              {
+                type: 'status_exception',
+                reason: 'Failed to find config with the provided config id: test_agent',
+              },
+            ],
+            type: 'status_exception',
+            reason: 'Failed to find config with the provided config id: test_agent',
+          },
+          status: 404,
+        },
+        statusCode: 404,
+      } as unknown) as ApiResponse)
+    );
+    await expect(
+      getAgentIdByConfig(client, 'test agent')
+    ).rejects.toThrowErrorMatchingInlineSnapshot(
+      `"Get agent 'test agent' failed, reason: {\\"error\\":{\\"root_cause\\":[{\\"type\\":\\"status_exception\\",\\"reason\\":\\"Failed to find config with the provided config id: test_agent\\"}],\\"type\\":\\"status_exception\\",\\"reason\\":\\"Failed to find config with the provided config id: test_agent\\"},\\"status\\":404}"`
+    );
+  });
+
+  it('handles search errors', async () => {
+    mockedTransport.mockRejectedValueOnce('request failed');
+    await expect(
+      getAgentIdByConfig(client, 'test agent')
+    ).rejects.toThrowErrorMatchingInlineSnapshot(
+      `"Get agent 'test agent' failed, reason: request failed"`
+    );
+  });
+
+  it('searches for agent id and sends request', async () => {
+    mockedTransport
+      .mockResolvedValueOnce({
+        body: {
+          type: 'agent',
+          configuration: { agent_id: 'new-id' },
+        },
+      })
+      .mockResolvedValueOnce({
+        body: { inference_results: [{ output: [{ result: 'test response' }] }] },
+      });
+    const response = await requestAgentByConfig({
+      context,
+      configName: 'new_agent',
+      body: { parameters: { param1: 'value1' } },
+    });
+    expect(mockedTransport).toBeCalledWith(
+      expect.objectContaining({ path: '/_plugins/_ml/agents/new-id/_execute' }),
+      expect.anything()
+    );
+    expect(response.body.inference_results[0].output[0].result).toEqual('test response');
+  });
+});