From c604baf19b62ef24404d07ca89d6809ba9286fc5 Mon Sep 17 00:00:00 2001
From: Liyun Xiu <xiliyun@amazon.com>
Date: Fri, 23 Aug 2024 08:54:59 +0800
Subject: [PATCH] Add UT

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
---
 .../server/routes/query_assist/agents.test.ts | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts b/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts
index 5e31191c9c97..a9ddbf16ea23 100644
--- a/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts
+++ b/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts
@@ -110,4 +110,27 @@ describe('Agents helper functions', () => {
     );
     expect(response.body.inference_results[0].output[0].result).toEqual('test response');
   });
+
+  it('searches for agent id and response contains ml_configuration', async () => {
+    mockedTransport
+      .mockResolvedValueOnce({
+        body: {
+          type: 'agent',
+          ml_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');
+  });
 });