Migrates kql_telemetry usage collector es client (#86585)

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>
elastic · Dec 21, 2020 · 61a0b00 · 61a0b00
1 parent 7de243e
commit 61a0b00
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 46 deletions.
diff --git a/src/plugins/data/server/kql_telemetry/usage_collector/fetch.test.ts b/src/plugins/data/server/kql_telemetry/usage_collector/fetch.test.ts
@@ -18,7 +18,7 @@
  */
 
 import { fetchProvider } from './fetch';
-import { LegacyAPICaller } from 'kibana/server';
+import { ElasticsearchClient } from 'kibana/server';
 import { CollectorFetchContext } from 'src/plugins/usage_collection/server';
 import { createCollectorFetchContextMock } from 'src/plugins/usage_collection/server/mocks';
 
@@ -30,42 +30,41 @@ jest.mock('../../../common', () => ({
 }));
 
 let fetch: ReturnType<typeof fetchProvider>;
-let callCluster: LegacyAPICaller;
+let esClient: ElasticsearchClient;
 let collectorFetchContext: CollectorFetchContext;
 const collectorFetchContextMock = createCollectorFetchContextMock();
 
 function setupMockCallCluster(
   optCount: { optInCount?: number; optOutCount?: number } | null,
   language: string | undefined | null
 ) {
-  callCluster = (jest.fn((method, params) => {
-    if (params && 'id' in params && params.id === 'kql-telemetry:kql-telemetry') {
-      if (optCount === null) {
-        return Promise.resolve({
+  function mockedEsGetMethod() {
+    if (optCount === null) {
+      return Promise.resolve({
+        body: {
           _index: '.kibana_1',
           _id: 'kql-telemetry:kql-telemetry',
           found: false,
-        });
-      } else {
-        return Promise.resolve({
+        },
+      });
+    } else {
+      return Promise.resolve({
+        body: {
           _source: {
-            'kql-telemetry': {
-              ...optCount,
-            },
+            'kql-telemetry': { ...optCount },
             type: 'kql-telemetry',
             updated_at: '2018-10-05T20:20:56.258Z',
           },
-        });
-      }
-    } else if (params && 'body' in params && params.body.query.term.type === 'config') {
-      if (language === 'missingConfigDoc') {
-        return Promise.resolve({
-          hits: {
-            hits: [],
-          },
-        });
-      } else {
-        return Promise.resolve({
+        },
+      });
+    }
+  }
+  function mockedEsSearchMethod() {
+    if (language === 'missingConfigDoc') {
+      return Promise.resolve({ body: { hits: { hits: [] } } });
+    } else {
+      return Promise.resolve({
+        body: {
           hits: {
             hits: [
               {
@@ -77,12 +76,15 @@ function setupMockCallCluster(
               },
             ],
           },
-        });
-      }
+        },
+      });
     }
-
-    throw new Error('invalid call');
-  }) as unknown) as LegacyAPICaller;
+  }
+  const esClientMock = ({
+    get: jest.fn().mockImplementation(mockedEsGetMethod),
+    search: jest.fn().mockImplementation(mockedEsSearchMethod),
+  } as unknown) as ElasticsearchClient;
+  esClient = esClientMock;
 }
 
 describe('makeKQLUsageCollector', () => {
@@ -95,7 +97,7 @@ describe('makeKQLUsageCollector', () => {
       setupMockCallCluster({ optInCount: 1 }, 'kuery');
       collectorFetchContext = {
         ...collectorFetchContextMock,
-        callCluster,
+        esClient,
       };
       const fetchResponse = await fetch(collectorFetchContext);
       expect(fetchResponse.optInCount).toBe(1);
@@ -106,7 +108,7 @@ describe('makeKQLUsageCollector', () => {
       setupMockCallCluster({ optInCount: 1 }, 'kuery');
       collectorFetchContext = {
         ...collectorFetchContextMock,
-        callCluster,
+        esClient,
       };
       const fetchResponse = await fetch(collectorFetchContext);
       expect(fetchResponse.defaultQueryLanguage).toBe('kuery');
@@ -117,7 +119,7 @@ describe('makeKQLUsageCollector', () => {
       setupMockCallCluster({ optInCount: 1 }, null);
       collectorFetchContext = {
         ...collectorFetchContextMock,
-        callCluster,
+        esClient,
       };
       const fetchResponse = await fetch(collectorFetchContext);
       expect(fetchResponse.defaultQueryLanguage).toBe('lucene');
@@ -127,7 +129,7 @@ describe('makeKQLUsageCollector', () => {
       setupMockCallCluster({ optInCount: 1 }, undefined);
       collectorFetchContext = {
         ...collectorFetchContextMock,
-        callCluster,
+        esClient,
       };
       const fetchResponse = await fetch(collectorFetchContext);
       expect(fetchResponse.defaultQueryLanguage).toBe('default-lucene');
@@ -137,7 +139,7 @@ describe('makeKQLUsageCollector', () => {
       setupMockCallCluster(null, 'kuery');
       collectorFetchContext = {
         ...collectorFetchContextMock,
-        callCluster,
+        esClient,
       };
       const fetchResponse = await fetch(collectorFetchContext);
       expect(fetchResponse.optInCount).toBe(0);
@@ -148,7 +150,7 @@ describe('makeKQLUsageCollector', () => {
       setupMockCallCluster(null, 'missingConfigDoc');
       collectorFetchContext = {
         ...collectorFetchContextMock,
-        callCluster,
+        esClient,
       };
       const fetchResponse = await fetch(collectorFetchContext);
       expect(fetchResponse.defaultQueryLanguage).toBe('default-lucene');

diff --git a/src/plugins/data/server/kql_telemetry/usage_collector/fetch.ts b/src/plugins/data/server/kql_telemetry/usage_collector/fetch.ts
@@ -30,18 +30,22 @@ export interface Usage {
 }
 
 export function fetchProvider(index: string) {
-  return async ({ callCluster }: CollectorFetchContext): Promise<Usage> => {
-    const [response, config] = await Promise.all([
-      callCluster('get', {
-        index,
-        id: 'kql-telemetry:kql-telemetry',
-        ignore: [404],
-      }),
-      callCluster('search', {
-        index,
-        body: { query: { term: { type: 'config' } } },
-        ignore: [404],
-      }),
+  return async ({ esClient }: CollectorFetchContext): Promise<Usage> => {
+    const [{ body: response }, { body: config }] = await Promise.all([
+      esClient.get(
+        {
+          index,
+          id: 'kql-telemetry:kql-telemetry',
+        },
+        { ignore: [404] }
+      ),
+      esClient.search(
+        {
+          index,
+          body: { query: { term: { type: 'config' } } },
+        },
+        { ignore: [404] }
+      ),
     ]);
 
     const queryLanguageConfigValue: string | null | undefined = get(