From 1eae08e2c1b1661bcc14561c0e141b47dc981ce5 Mon Sep 17 00:00:00 2001
From: Michael Dokolin <mikhail.dokolin@elastic.co>
Date: Wed, 11 Aug 2021 22:16:53 +0200
Subject: [PATCH] [Expressions] Add support of partial results to the switch
 expression function (#108086)

---
 .../functions/common/switch.test.js           | 68 +++++++++++++------
 .../functions/common/switch.ts                | 19 +++---
 2 files changed, 59 insertions(+), 28 deletions(-)

diff --git a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.test.js b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.test.js
index ffa1557d2b54..d0ea8fff0a45 100644
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.test.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.test.js
@@ -56,25 +56,6 @@ describe('switch', () => {
   });
 
   describe('function', () => {
-    describe('with no cases', () => {
-      it('should return the context if no default is provided', () => {
-        const context = 'foo';
-
-        testScheduler.run(({ expectObservable }) =>
-          expectObservable(fn(context, {})).toBe('(0|)', [context])
-        );
-      });
-
-      it('should return the default if provided', () => {
-        const context = 'foo';
-        const args = { default: () => of('bar') };
-
-        testScheduler.run(({ expectObservable }) =>
-          expectObservable(fn(context, args)).toBe('(0|)', ['bar'])
-        );
-      });
-    });
-
     describe('with no matching cases', () => {
       it('should return the context if no default is provided', () => {
         const context = 'foo';
@@ -108,6 +89,55 @@ describe('switch', () => {
           expectObservable(fn(context, args)).toBe('(0|)', [result])
         );
       });
+
+      it('should support partial results', () => {
+        testScheduler.run(({ cold, expectObservable }) => {
+          const context = 'foo';
+          const case1 = cold('--ab-c-', {
+            a: {
+              type: 'case',
+              matches: false,
+              result: 1,
+            },
+            b: {
+              type: 'case',
+              matches: true,
+              result: 2,
+            },
+            c: {
+              type: 'case',
+              matches: false,
+              result: 3,
+            },
+          });
+          const case2 = cold('-a--bc-', {
+            a: {
+              type: 'case',
+              matches: true,
+              result: 4,
+            },
+            b: {
+              type: 'case',
+              matches: true,
+              result: 5,
+            },
+            c: {
+              type: 'case',
+              matches: true,
+              result: 6,
+            },
+          });
+          const expected = ' --abc(de)-';
+          const args = { case: [() => case1, () => case2] };
+          expectObservable(fn(context, args)).toBe(expected, {
+            a: 4,
+            b: 2,
+            c: 2,
+            d: 5,
+            e: 6,
+          });
+        });
+      });
     });
   });
 });
diff --git a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.ts b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.ts
index f4e6c92c91cb..3e676c829a30 100644
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.ts
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.ts
@@ -5,14 +5,14 @@
  * 2.0.
  */
 
-import { Observable, defer, from, of } from 'rxjs';
-import { concatMap, filter, merge, pluck, take } from 'rxjs/operators';
+import { Observable, combineLatest, defer, of } from 'rxjs';
+import { concatMap } from 'rxjs/operators';
 import { ExpressionFunctionDefinition } from 'src/plugins/expressions/common';
 import { Case } from '../../../types';
 import { getFunctionHelp } from '../../../i18n';
 
 interface Arguments {
-  case?: Array<() => Observable<Case>>;
+  case: Array<() => Observable<Case>>;
   default?(): Observable<any>;
 }
 
@@ -43,12 +43,13 @@ export function switchFn(): ExpressionFunctionDefinition<
       },
     },
     fn(input, args) {
-      return from(args.case ?? []).pipe(
-        concatMap((item) => item()),
-        filter(({ matches }) => matches),
-        pluck('result'),
-        merge(defer(() => args.default?.() ?? of(input))),
-        take(1)
+      return combineLatest(args.case.map((item) => defer(() => item()))).pipe(
+        concatMap((items) => {
+          const item = items.find(({ matches }) => matches);
+          const item$ = item && of(item.result);
+
+          return item$ ?? args.default?.() ?? of(input);
+        })
       );
     },
   };