catalog: fix EntitySwitch unmounting children on entity refresh

Signed-off-by: Vincenzo Scamporlino <vincenzos@spotify.com>

plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx

@@ -20,7 +20,7 @@ import {
   EntityProvider,
 } from '@backstage/plugin-catalog-react';
 import { render, screen } from '@testing-library/react';
-import React from 'react';
+import React, { useEffect } from 'react';
 import { isKind } from './conditions';
 import { EntitySwitch } from './EntitySwitch';
 import { featureFlagsApiRef } from '@backstage/core-plugin-api';
@@ -209,6 +209,97 @@ describe('EntitySwitch', () => {
     expect(screen.queryByText('B')).not.toBeInTheDocument();
   });
 
+  it('should display the children in case entity is available and loading is true', () => {
+    const entity = { metadata: { name: 'mock' }, kind: 'component' } as Entity;
+
+    render(
+      <Wrapper>
+        <AsyncEntityProvider entity={entity} loading>
+          <EntitySwitch>
+            <EntitySwitch.Case
+              if={isKind('component')}
+              children={<p>Component</p>}
+            />
+            <EntitySwitch.Case if={isKind('system')} children={<p>System</p>} />
+          </EntitySwitch>
+        </AsyncEntityProvider>
+      </Wrapper>,
+    );
+
+    expect(screen.queryByText('Component')).toBeInTheDocument();
+    expect(screen.queryByText('System')).not.toBeInTheDocument();
+  });
+
+  it(`shouldn't unmount the children on entity refresh`, () => {
+    const entity = { metadata: { name: 'mock' }, kind: 'component' } as Entity;
+
+    let mountsCount = 0;
+    function Component() {
+      useEffect(() => {
+        ++mountsCount;
+      }, []);
+
+      return <p>Component</p>;
+    }
+
+    const rendered = render(
+      <Wrapper>
+        <AsyncEntityProvider entity={entity} loading={false}>
+          <EntitySwitch>
+            <EntitySwitch.Case
+              if={isKind('component')}
+              children={<Component />}
+            />
+            <EntitySwitch.Case if={isKind('system')} children={<p>System</p>} />
+          </EntitySwitch>
+        </AsyncEntityProvider>
+      </Wrapper>,
+    );
+
+    expect(screen.queryByText('Component')).toBeInTheDocument();
+    expect(screen.queryByText('System')).not.toBeInTheDocument();
+
+    expect(mountsCount).toBe(1);
+
+    rendered.rerender(
+      <Wrapper>
+        <AsyncEntityProvider entity={entity} loading>
+          <EntitySwitch>
+            <EntitySwitch.Case
+              if={isKind('component')}
+              children={<Component />}
+            />
+            <EntitySwitch.Case if={isKind('system')} children={<p>System</p>} />
+          </EntitySwitch>
+        </AsyncEntityProvider>
+      </Wrapper>,
+    );
+
+    expect(screen.queryByText('Component')).toBeInTheDocument();
+    expect(screen.queryByText('System')).not.toBeInTheDocument();
+
+    expect(mountsCount).toBe(1);
+
+    rendered.rerender(
+      <Wrapper>
+        <AsyncEntityProvider entity={entity} loading={false}>
+          <EntitySwitch>
+            <EntitySwitch.Case
+              if={isKind('component')}
+              children={<Component />}
+            />
+            <EntitySwitch.Case if={isKind('system')} children={<p>System</p>} />
+          </EntitySwitch>
+        </AsyncEntityProvider>
+      </Wrapper>,
+    );
+
+    expect(screen.queryByText('Component')).toBeInTheDocument();
+    expect(screen.queryByText('System')).not.toBeInTheDocument();
+
+    expect(mountsCount).toBe(1);
+  });
+
   it('should switch with async condition that is true', async () => {
     const entity = { metadata: { name: 'mock' }, kind: 'component' } as Entity;
 

plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx

@@ -64,8 +64,9 @@ export interface EntitySwitchProps {
 
 /** @public */
 export const EntitySwitch = (props: EntitySwitchProps) => {
-  const { entity, loading } = useAsyncEntity();
+  const { entity } = useAsyncEntity();
   const apis = useApiHolder();
+
   const results = useElementFilter(
     props.children,
     collection =>
@@ -76,11 +77,6 @@ export const EntitySwitch = (props: EntitySwitchProps) => {
         })
         .getElements()
         .flatMap<SwitchCaseResult>((element: ReactElement) => {
-          // Nothing is rendered while loading
-          if (loading) {
-            return [];
-          }
-
           const { if: condition, children: elementsChildren } =
             element.props as EntitySwitchCase;
 
@@ -100,7 +96,7 @@ export const EntitySwitch = (props: EntitySwitchProps) => {
             },
           ];
         }),
-    [apis, entity, loading],
+    [apis, entity],
   );
 
   const hasAsyncCases = results.some(