From 4e3755151c358ae04f59408c7bcaec3bdfdf6ffd Mon Sep 17 00:00:00 2001
From: Johan Haals <johan.haals@gmail.com>
Date: Tue, 5 Oct 2021 14:41:24 +0200
Subject: [PATCH] Keep existing EntitySwitch behaviour, add tests

Keep the existing behavior for non async statements

Co-authored-by: Patrik Oldsberg <poldsberg@gmail.com>
Signed-off-by: Johan Haals <johan.haals@gmail.com>
---
 .../EntitySwitch/EntitySwitch.test.tsx        | 59 ++++++++++++++++
 .../components/EntitySwitch/EntitySwitch.tsx  | 70 +++++++++++++------
 2 files changed, 106 insertions(+), 23 deletions(-)

diff --git a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx
index f80a66b9bb..7f8bf35b74 100644
--- a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx
+++ b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx
@@ -112,4 +112,63 @@ describe('EntitySwitch', () => {
     expect(rendered.queryByText('A')).not.toBeInTheDocument();
     expect(rendered.queryByText('B')).toBeInTheDocument();
   });
+
+  it('should switch with async condition that is true', async () => {
+    const entity = { kind: 'component' } as Entity;
+
+    const shouldRender = () => Promise.resolve(true);
+    const rendered = render(
+      <Wrapper>
+        <EntityProvider entity={entity}>
+          <EntitySwitch>
+            <EntitySwitch.Case if={shouldRender} children="A" />
+            <EntitySwitch.Case children="B" />
+          </EntitySwitch>
+        </EntityProvider>
+      </Wrapper>,
+    );
+
+    await expect(rendered.findByText('A')).resolves.toBeInTheDocument();
+    expect(rendered.queryByText('B')).not.toBeInTheDocument();
+  });
+
+  it('should switch with sync condition that is false', async () => {
+    const entity = { kind: 'component' } as Entity;
+
+    const shouldRender = () => Promise.resolve(false);
+    const rendered = render(
+      <Wrapper>
+        <EntityProvider entity={entity}>
+          <EntitySwitch>
+            <EntitySwitch.Case if={shouldRender} children="A" />
+            <EntitySwitch.Case if={() => true} children="B" />
+          </EntitySwitch>
+        </EntityProvider>
+      </Wrapper>,
+    );
+
+    await expect(rendered.findByText('B')).resolves.toBeInTheDocument();
+    expect(rendered.queryByText('A')).not.toBeInTheDocument();
+  });
+
+  it('should switch with sync condition that throws', async () => {
+    const entity = { kind: 'component' } as Entity;
+
+    const shouldRender = () => Promise.reject();
+    const rendered = render(
+      <Wrapper>
+        <EntityProvider entity={entity}>
+          <EntitySwitch>
+            <EntitySwitch.Case if={shouldRender} children="A" />
+            <EntitySwitch.Case if={() => false} children="B" />
+            <EntitySwitch.Case children="C" />
+          </EntitySwitch>
+        </EntityProvider>
+      </Wrapper>,
+    );
+
+    await expect(rendered.findByText('C')).resolves.toBeInTheDocument();
+    expect(rendered.queryByText('A')).not.toBeInTheDocument();
+    expect(rendered.queryByText('B')).not.toBeInTheDocument();
+  });
 });
diff --git a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx
index bf2caac525..5c50c2507c 100644
--- a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx
+++ b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx
@@ -16,7 +16,7 @@
 
 import { Entity } from '@backstage/catalog-model';
 import { useEntity } from '@backstage/plugin-catalog-react';
-import { PropsWithChildren, ReactNode } from 'react';
+import React, { PropsWithChildren, ReactNode } from 'react';
 import {
   attachComponentData,
   useApiHolder,
@@ -45,45 +45,69 @@ type SwitchCase = {
   children: JSX.Element;
 };
 
+type SwitchCaseResult = {
+  if: boolean | Promise<boolean>;
+  children: JSX.Element;
+};
+
 export const EntitySwitch = ({ children }: PropsWithChildren<{}>) => {
   const { entity } = useEntity();
   const apis = useApiHolder();
-  const switchCases = useElementFilter(children, collection =>
-    collection
-      .selectByComponentData({
-        key: ENTITY_SWITCH_KEY,
-        withStrictError: 'Child of EntitySwitch is not an EntitySwitch.Case',
-      })
-      .getElements()
-      .flatMap<SwitchCase>((element: React.ReactElement) => {
-        const { if: condition, children: elementsChildren } = element.props;
-        return [{ if: condition, children: elementsChildren }];
-      }),
+  const results = useElementFilter(
+    children,
+    collection =>
+      collection
+        .selectByComponentData({
+          key: ENTITY_SWITCH_KEY,
+          withStrictError: 'Child of EntitySwitch is not an EntitySwitch.Case',
+        })
+        .getElements()
+        .flatMap<SwitchCaseResult>((element: React.ReactElement) => {
+          const { if: condition, children: elementsChildren } =
+            element.props as SwitchCase;
+          return [
+            {
+              if: condition?.(entity, { apis }) ?? true,
+              children: elementsChildren,
+            },
+          ];
+        }),
+    [apis, entity],
+  );
+  const hasAsyncCases = results.some(
+    r => typeof r.if === 'object' && 'then' in r.if,
   );
 
+  if (hasAsyncCases) {
+    return <AsyncEntitySwitch results={results} />;
+  }
+
+  return results.find(r => r.if)?.children ?? null;
+};
+
+function AsyncEntitySwitch({ results }: { results: SwitchCaseResult[] }) {
   const { loading, value } = useAsync(async () => {
-    const promises = switchCases.map(
+    const promises = results.map(
       async ({ if: condition, children: output }) => {
-        if (!condition) {
-          return output;
+        try {
+          if (await condition) {
+            return output;
+          }
+        } catch {
+          /* ignored */
         }
 
-        try {
-          const matches = await condition(entity, { apis });
-          return matches ? output : null;
-        } catch {
-          return null;
-        }
+        return null;
       },
     );
     return (await Promise.all(promises)).find(Boolean) ?? null;
-  }, [switchCases]);
+  }, [results]);
 
   if (loading || !value) {
     return null;
   }
 
   return value;
-};
+}
 
 EntitySwitch.Case = EntitySwitchCase;