diff --git a/plugins/catalog/package.json b/plugins/catalog/package.json
index e1e3b388a1..16e756eabe 100644
--- a/plugins/catalog/package.json
+++ b/plugins/catalog/package.json
@@ -55,6 +55,7 @@
   },
   "devDependencies": {
     "@backstage/cli": "^0.7.0",
+    "@backstage/core-app-api": "^0.1.2",
     "@backstage/dev-utils": "^0.1.17",
     "@backstage/test-utils": "^0.1.13",
     "@testing-library/jest-dom": "^5.10.1",
diff --git a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx
index 20b9aaba4f..5b21266d44 100644
--- a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx
+++ b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.test.tsx
@@ -20,6 +20,19 @@ import { render } from '@testing-library/react';
 import React from 'react';
 import { isKind } from './conditions';
 import { EntitySwitch } from './EntitySwitch';
+import { featureFlagsApiRef } from '@backstage/core-plugin-api';
+import {
+  LocalStorageFeatureFlags,
+  ApiProvider,
+  ApiRegistry,
+} from '@backstage/core-app-api';
+
+const mockFeatureFlagsApi = new LocalStorageFeatureFlags();
+const Wrapper = ({ children }: { children?: React.ReactNode }) => (
+  <ApiProvider apis={ApiRegistry.with(featureFlagsApiRef, mockFeatureFlagsApi)}>
+    {children}
+  </ApiProvider>
+);
 
 describe('EntitySwitch', () => {
   it('should switch child when entity switches', () => {
@@ -32,15 +45,17 @@ describe('EntitySwitch', () => {
     );
 
     const rendered = render(
-      <EntityContext.Provider
-        value={{
-          entity: { kind: 'component' } as Entity,
-          loading: false,
-          error: undefined,
-        }}
-      >
-        {content}
-      </EntityContext.Provider>,
+      <Wrapper>
+        <EntityContext.Provider
+          value={{
+            entity: { kind: 'component' } as Entity,
+            loading: false,
+            error: undefined,
+          }}
+        >
+          {content}
+        </EntityContext.Provider>
+      </Wrapper>,
     );
 
     expect(rendered.queryByText('A')).toBeInTheDocument();
@@ -48,15 +63,17 @@ describe('EntitySwitch', () => {
     expect(rendered.queryByText('C')).not.toBeInTheDocument();
 
     rendered.rerender(
-      <EntityContext.Provider
-        value={{
-          entity: { kind: 'template' } as Entity,
-          loading: false,
-          error: undefined,
-        }}
-      >
-        {content}
-      </EntityContext.Provider>,
+      <Wrapper>
+        <EntityContext.Provider
+          value={{
+            entity: { kind: 'template' } as Entity,
+            loading: false,
+            error: undefined,
+          }}
+        >
+          {content}
+        </EntityContext.Provider>
+      </Wrapper>,
     );
 
     expect(rendered.queryByText('A')).not.toBeInTheDocument();
@@ -64,15 +81,17 @@ describe('EntitySwitch', () => {
     expect(rendered.queryByText('C')).not.toBeInTheDocument();
 
     rendered.rerender(
-      <EntityContext.Provider
-        value={{
-          entity: { kind: 'derp' } as Entity,
-          loading: false,
-          error: undefined,
-        }}
-      >
-        {content}
-      </EntityContext.Provider>,
+      <Wrapper>
+        <EntityContext.Provider
+          value={{
+            entity: { kind: 'derp' } as Entity,
+            loading: false,
+            error: undefined,
+          }}
+        >
+          {content}
+        </EntityContext.Provider>
+      </Wrapper>,
     );
 
     expect(rendered.queryByText('A')).not.toBeInTheDocument();
@@ -88,24 +107,28 @@ describe('EntitySwitch', () => {
     };
 
     const rendered = render(
-      <EntityContext.Provider value={entityContextValue}>
-        <EntitySwitch>
-          <EntitySwitch.Case if={isKind('component')} children="A" />
-          <EntitySwitch.Case children="B" />
-        </EntitySwitch>
-      </EntityContext.Provider>,
+      <Wrapper>
+        <EntityContext.Provider value={entityContextValue}>
+          <EntitySwitch>
+            <EntitySwitch.Case if={isKind('component')} children="A" />
+            <EntitySwitch.Case children="B" />
+          </EntitySwitch>
+        </EntityContext.Provider>
+      </Wrapper>,
     );
 
     expect(rendered.queryByText('A')).toBeInTheDocument();
     expect(rendered.queryByText('B')).not.toBeInTheDocument();
 
     rendered.rerender(
-      <EntityContext.Provider value={entityContextValue}>
-        <EntitySwitch>
-          <EntitySwitch.Case if={isKind('template')} children="A" />
-          <EntitySwitch.Case children="B" />
-        </EntitySwitch>
-      </EntityContext.Provider>,
+      <Wrapper>
+        <EntityContext.Provider value={entityContextValue}>
+          <EntitySwitch>
+            <EntitySwitch.Case if={isKind('template')} children="A" />
+            <EntitySwitch.Case children="B" />
+          </EntitySwitch>
+        </EntityContext.Provider>
+      </Wrapper>,
     );
 
     expect(rendered.queryByText('A')).not.toBeInTheDocument();
diff --git a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx
index a3ede42288..4b12a7249b 100644
--- a/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx
+++ b/plugins/catalog/src/components/EntitySwitch/EntitySwitch.tsx
@@ -16,49 +16,40 @@
 
 import { Entity } from '@backstage/catalog-model';
 import { useEntity } from '@backstage/plugin-catalog-react';
+import { PropsWithChildren, ReactNode } from 'react';
 import {
-  Children,
-  Fragment,
-  isValidElement,
-  PropsWithChildren,
-  ReactNode,
-  useMemo,
-} from 'react';
+  attachComponentData,
+  useElementFilter,
+} from '@backstage/core-plugin-api';
+
+const ENTITY_SWITCH_KEY = 'core.backstage.entitySwitch';
 
 const EntitySwitchCase = (_: {
   if?: (entity: Entity) => boolean;
   children: ReactNode;
 }) => null;
 
+attachComponentData(EntitySwitchCase, ENTITY_SWITCH_KEY, true);
+
 type SwitchCase = {
   if?: (entity: Entity) => boolean;
   children: JSX.Element;
 };
 
-function createSwitchCasesFromChildren(childrenNode: ReactNode): SwitchCase[] {
-  return Children.toArray(childrenNode).flatMap(child => {
-    if (!isValidElement(child)) {
-      return [];
-    }
-
-    if (child.type === Fragment) {
-      return createSwitchCasesFromChildren(child.props.children);
-    }
-
-    if (child.type !== EntitySwitchCase) {
-      throw new Error(`Child of EntitySwitch is not an EntitySwitch.Case`);
-    }
-
-    const { if: condition, children } = child.props;
-    return [{ if: condition, children }];
-  });
-}
-
 export const EntitySwitch = ({ children }: PropsWithChildren<{}>) => {
   const { entity } = useEntity();
-  const switchCases = useMemo(() => createSwitchCasesFromChildren(children), [
-    children,
-  ]);
+  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 matchingCase = switchCases.find(switchCase =>
     switchCase.if ? switchCase.if(entity) : true,