Add functional test to ensure painless script is working

x-pack/legacy/plugins/lens/server/usage/visualization_counts.ts

@@ -5,11 +5,20 @@
  */
 
 import { KibanaConfig } from 'src/legacy/server/kbn_server';
-import { CallCluster } from 'src/legacy/core_plugins/elasticsearch';
+import { CallClusterOptions } from 'src/legacy/core_plugins/elasticsearch';
+import { SearchParams, SearchResponse } from 'elasticsearch';
 import { LensUsage } from './types';
 
+type ClusterSearchType = (
+  endpoint: 'search',
+  params: SearchParams & {
+    rest_total_hits_as_int: boolean;
+  },
+  options?: CallClusterOptions
+) => Promise<SearchResponse<unknown>>;
+
 export async function getVisualizationCounts(
-  callCluster: CallCluster,
+  callCluster: ClusterSearchType,
   config: KibanaConfig
 ): Promise<LensUsage> {
   const scriptedMetric = {
@@ -28,11 +37,14 @@ export async function getVisualizationCounts(
       // Despite docs that say this is optional, this script can't be blank.
       combine_script: 'return state',
       // Reduce script is executed across all clusters, so we need to add up all the total from each cluster
+      // This also needs to account for having no data
       reduce_script: `
         Map result = [:];
         for (Map m : states.toArray()) {
-          for (String k : m.keySet()) {
-            result.put(k, result.containsKey(k) ? result.get(k) + m.get(k) : m.get(k));
+          if (m !== null) {
+            for (String k : m.keySet()) {
+              result.put(k, result.containsKey(k) ? result.get(k) + m.get(k) : m.get(k));
+            }
           }
         }
         return result;

x-pack/test/functional/apps/lens/smokescreen.ts

@@ -4,11 +4,15 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+import _ from 'lodash';
 import expect from '@kbn/expect';
+import { Client, SearchParams } from 'elasticsearch';
+import { KibanaConfig } from 'src/legacy/server/kbn_server';
 import { FtrProviderContext } from '../../ftr_provider_context';
+import { getVisualizationCounts } from '../../../../legacy/plugins/lens/server/usage/visualization_counts';
 
 // eslint-disable-next-line import/no-default-export
-export default function({ getService, getPageObjects }: FtrProviderContext) {
+export default function({ getService, getPageObjects, ...rest }: FtrProviderContext) {
   const PageObjects = getPageObjects([
     'header',
     'common',
@@ -112,5 +116,34 @@ export default function({ getService, getPageObjects }: FtrProviderContext) {
       // legend item(s), so we're using a class selector here.
       expect(await find.allByCssSelector('.echLegendItem')).to.have.length(3);
     });
+
+    it('should collect telemetry on saved visualization types with a painless script', async () => {
+      const es: Client = getService('es');
+      const callCluster = (path: 'search', searchParams: SearchParams) =>
+        es[path].call(es, searchParams);
+
+      const results = await getVisualizationCounts(callCluster, {
+        // Fake KibanaConfig service
+        get(key: string) {
+          return '.kibana';
+        },
+        has: () => false,
+      } as KibanaConfig);
+
+      expect(results).to.have.keys([
+        'visualization_types_overall',
+        'visualization_types_last_30_days',
+        'visualization_types_last_90_days',
+        'saved_total',
+        'saved_last_30_days',
+        'saved_last_90_days',
+      ]);
+
+      expect(results.visualization_types_overall).to.eql({
+        lnsMetric: 1,
+        bar_stacked: 1,
+      });
+      expect(results.saved_total).to.eql(2);
+    });
   });
 }