feat: more work with the nice API

Co-authored-by: Patrik Oldsberg <poldsberg@gmail.com>
Co-authored-by: Fredrik Adelöw <freben@users.noreply.github.com>
Co-authored-by: Johan Haals <johan.haals@gmail.com>

Signed-off-by: blam <ben@blam.sh>

plugins/catalog/src/components/EntityLayout/EntityLayout.tsx

@@ -28,6 +28,10 @@ import {
   Page,
   Progress,
   RoutedTabs,
+  FeatureFlagsApi,
+  getComponentData,
+  featureFlagsApiRef,
+  useApi,
 } from '@backstage/core';
 import {
   EntityContext,
@@ -58,46 +62,14 @@ type SubRoute = {
   tabProps?: TabProps<React.ElementType, { component?: React.ElementType }>;
 };
 
+const dataKey = 'plugin.catalog.entityLayoutRoute';
+
 const Route: (props: SubRoute) => null = () => null;
+attachComponentData(Route, dataKey, true);
 
 // This causes all mount points that are discovered within this route to use the path of the route itself
 attachComponentData(Route, 'core.gatherMountPoints', true);
 
-function createSubRoutesFromChildren(
-  childrenProps: React.ReactNode,
-  entity: Entity | undefined,
-): SubRoute[] {
-  // Directly comparing child.type with Route will not work with in
-  // combination with react-hot-loader in storybook
-  // https://github.com/gaearon/react-hot-loader/issues/304
-  const routeType = (
-    <Route path="" title="">
-      <div />
-    </Route>
-  ).type;
-
-  return Children.toArray(childrenProps).flatMap(child => {
-    if (!isValidElement(child)) {
-      return [];
-    }
-
-    if (child.type === Fragment) {
-      return createSubRoutesFromChildren(child.props.children, entity);
-    }
-
-    if (child.type !== routeType) {
-      throw new Error('Child of EntityLayout must be an EntityLayout.Route');
-    }
-
-    const { path, title, children, if: condition, tabProps } = child.props;
-    if (condition && entity && !condition(entity)) {
-      return [];
-    }
-
-    return [{ path, title, children, tabProps }];
-  });
-}
-
 const EntityLayoutTitle = ({
   entity,
   title,
@@ -195,7 +167,27 @@ export const EntityLayout = ({
   const { kind, namespace, name } = useEntityCompoundName();
   const { entity, loading, error } = useContext(EntityContext);
 
-  const routes = createSubRoutesFromChildren(children, entity);
+  const routes = useElementCollection(children)
+    .findByComponentData({
+      key: dataKey,
+      withStrictError: 'Child of EntityLayout must be an EntityLayout.Route',
+    })
+    .listElements() // all nodes, element data, maintain structure or not?
+    .flatMap(({ props }) => {
+      if (props.condition && entity && !props.condition(entity)) {
+        return [];
+      }
+
+      return [
+        {
+          path: props.path,
+          title: props.title,
+          children: props.children,
+          tabProps: props.tabProps,
+        },
+      ];
+    });
+
   const { headerTitle, headerType } = headerProps(
     kind,
     namespace,