From 21ad2d858a7a3a22515816605506dd47c0b85b11 Mon Sep 17 00:00:00 2001
From: Xinrui Bai <xinruiba@amazon.com>
Date: Mon, 26 Feb 2024 23:05:27 +0000
Subject: [PATCH] [UT] Update unit test

Signed-off-by: Xinrui Bai <xinruiba@amazon.com>
---
 .../opensearch_search_strategy.test.ts        | 50 +++++++++++++++++--
 .../opensearch_search_strategy.ts             |  2 +-
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts b/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts
index 8ad048de35e6..ae6e1746dab6 100644
--- a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts
+++ b/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts
@@ -56,6 +56,19 @@ describe('OpenSearch search strategy', () => {
     },
     dataSourceEnabled: jest.fn(() => true),
     registerCredentialProvider: jest.fn(),
+    registerCustomApiSchema(schema: any): void {
+      throw new Error('Function not implemented.');
+    },
+  };
+  const mockDataSourcePluginSetupWithDataSourceDisabled: DataSourcePluginSetup = {
+    createDataSourceError(err: any): DataSourceError {
+      return new DataSourceError({});
+    },
+    dataSourceEnabled: jest.fn(() => false),
+    registerCredentialProvider: jest.fn(),
+    registerCustomApiSchema(schema: any): void {
+      throw new Error('Function not implemented.');
+    },
   };
   const body = {
     body: {
@@ -69,6 +82,7 @@ describe('OpenSearch search strategy', () => {
   };
   const mockOpenSearchApiCaller = jest.fn().mockResolvedValue(body);
   const mockDataSourceApiCaller = jest.fn().mockResolvedValue(body);
+  const mockOpenSearchApiCallerWithLongNumeralsSupport = jest.fn().mockResolvedValue(body);
   const dataSourceId = 'test-data-source-id';
   const mockDataSourceContext = {
     dataSource: {
@@ -86,7 +100,14 @@ describe('OpenSearch search strategy', () => {
           get: () => {},
         },
       },
-      opensearch: { client: { asCurrentUser: { search: mockOpenSearchApiCaller } } },
+      opensearch: {
+        client: {
+          asCurrentUser: { search: mockOpenSearchApiCaller },
+          asCurrentUserWithLongNumeralsSupport: {
+            search: mockOpenSearchApiCallerWithLongNumeralsSupport,
+          },
+        },
+      },
     },
   };
   const mockDataSourceEnabledContext = {
@@ -215,10 +236,10 @@ describe('OpenSearch search strategy', () => {
     const hostsTobeTested = [undefined, []];
     const dataSourceIdToBeTested = [undefined, '', dataSourceId];
 
-    hostsTobeTested.forEach(() => {
+    hostsTobeTested.forEach((host) => {
       const mockOpenSearchServiceSetup = opensearchServiceMock.createSetup();
 
-      if (hostsTobeTested !== undefined) {
+      if (host !== undefined) {
         mockOpenSearchServiceSetup.legacy.client = {
           callAsInternalUser: jest.fn(),
           asScoped: jest.fn(),
@@ -276,4 +297,27 @@ describe('OpenSearch search strategy', () => {
       expect(mockDataSourceApiCaller).not.toBeCalled();
     }
   });
+
+  it('dataSource disabled and longNumeralsSupported, send request without dataSourceId should get longNumeralsSupport client', async () => {
+    const mockOpenSearchServiceSetup = opensearchServiceMock.createSetup();
+    const opensearchSearch = await opensearchSearchStrategyProvider(
+      mockConfig$,
+      mockLogger,
+      mockSearchUsage,
+      mockDataSourcePluginSetupWithDataSourceDisabled,
+      mockOpenSearchServiceSetup,
+      true
+    );
+
+    const dataSourceIdToBeTested = [undefined, ''];
+
+    for (const testDataSourceId of dataSourceIdToBeTested) {
+      await opensearchSearch.search((mockContext as unknown) as RequestHandlerContext, {
+        dataSourceId: testDataSourceId,
+      });
+      expect(mockOpenSearchApiCallerWithLongNumeralsSupport).toBeCalled();
+      expect(mockOpenSearchApiCaller).not.toBeCalled();
+      expect(mockDataSourceApiCaller).not.toBeCalled();
+    }
+  });
 });
diff --git a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts b/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts
index 0e77ea20a6b1..fa1b3e4da94c 100644
--- a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts
+++ b/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts
@@ -75,7 +75,7 @@ export const opensearchSearchStrategyProvider = (
 
       try {
         const isOpenSearchHostsEmpty =
-          openSearchServiceSetup?.legacy.client.config.hosts.length === 0;
+          openSearchServiceSetup?.legacy?.client?.config?.hosts?.length === 0;
 
         if (dataSource?.dataSourceEnabled() && isOpenSearchHostsEmpty && !request.dataSourceId) {
           throw new Error(`Data source id is required when no openseach hosts config provided`);