From 4491ba231fe55a3913962200cad3f7d39205c321 Mon Sep 17 00:00:00 2001
From: Adam Harvey <adam.harvey@dxc.com>
Date: Thu, 25 Mar 2021 09:50:35 -0400
Subject: [PATCH] Refactor entity collection

Signed-off-by: Adam Harvey <adam.harvey@dxc.com>
---
 .changeset/perfect-coats-train.md                  |  2 +-
 packages/app/src/components/catalog/EntityPage.tsx |  2 +-
 .../SystemDiagramCard/SystemDiagramCard.test.tsx   | 14 +++++++++++---
 .../SystemDiagramCard/SystemDiagramCard.tsx        |  8 +++-----
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/.changeset/perfect-coats-train.md b/.changeset/perfect-coats-train.md
index 98b558e8d6..7a089e870f 100644
--- a/.changeset/perfect-coats-train.md
+++ b/.changeset/perfect-coats-train.md
@@ -19,7 +19,7 @@ For example,
 +    <EntityPageLayout.Content
 +      path="/diagram/*"
 +      title="Diagram"
-+      element={<EntitySystemDiagramCard entity={entity} />}
++      element={<EntitySystemDiagramCard />}
 +    />
    </EntityPageLayout>
  );
diff --git a/packages/app/src/components/catalog/EntityPage.tsx b/packages/app/src/components/catalog/EntityPage.tsx
index 2d0bde2a7a..e46186bbe0 100644
--- a/packages/app/src/components/catalog/EntityPage.tsx
+++ b/packages/app/src/components/catalog/EntityPage.tsx
@@ -461,7 +461,7 @@ const SystemEntityPage = ({ entity }: { entity: Entity }) => (
     <EntityPageLayout.Content
       path="/diagram/*"
       title="Diagram"
-      element={<EntitySystemDiagramCard entity={entity} />}
+      element={<EntitySystemDiagramCard />}
     />
   </EntityPageLayout>
 );
diff --git a/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.test.tsx b/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.test.tsx
index ed1a1a69ea..baa8c667d2 100644
--- a/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.test.tsx
+++ b/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.test.tsx
@@ -15,7 +15,11 @@
  */
 
 import { ApiProvider, ApiRegistry } from '@backstage/core';
-import { catalogApiRef, CatalogApi } from '@backstage/plugin-catalog-react';
+import {
+  catalogApiRef,
+  CatalogApi,
+  EntityProvider,
+} from '@backstage/plugin-catalog-react';
 import { Entity, RELATION_PART_OF } from '@backstage/catalog-model';
 import { renderInTestApp } from '@backstage/test-utils';
 import React from 'react';
@@ -51,7 +55,9 @@ describe('<SystemDiagramCard />', () => {
 
     const { queryByText } = await renderInTestApp(
       <ApiProvider apis={ApiRegistry.from([[catalogApiRef, catalogApi]])}>
-        <SystemDiagramCard entity={entity} />
+        <EntityProvider entity={entity}>
+          <SystemDiagramCard />
+        </EntityProvider>
       </ApiProvider>,
     );
 
@@ -105,7 +111,9 @@ describe('<SystemDiagramCard />', () => {
 
     const { getByText } = await renderInTestApp(
       <ApiProvider apis={ApiRegistry.from([[catalogApiRef, catalogApi]])}>
-        <SystemDiagramCard entity={entity} />
+        <EntityProvider entity={entity}>
+          <SystemDiagramCard />
+        </EntityProvider>
       </ApiProvider>,
     );
 
diff --git a/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.tsx b/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.tsx
index bec2e844c8..698476fe93 100644
--- a/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.tsx
+++ b/plugins/catalog/src/components/SystemDiagramCard/SystemDiagramCard.tsx
@@ -23,6 +23,7 @@ import {
 import {
   catalogApiRef,
   getEntityRelations,
+  useEntity,
 } from '@backstage/plugin-catalog-react';
 import {
   DependencyGraph,
@@ -37,10 +38,6 @@ import ZoomOutMap from '@material-ui/icons/ZoomOutMap';
 import React from 'react';
 import { useAsync } from 'react-use';
 
-type SystemDiagramProps = {
-  entity: Entity;
-};
-
 function simplifiedEntityName(
   ref:
     | Entity
@@ -61,7 +58,8 @@ function simplifiedEntityName(
  * Dynamically generates a diagram of a system, its assigned entities,
  * and relationships of those entities.
  */
-export function SystemDiagramCard({ entity }: SystemDiagramProps) {
+export function SystemDiagramCard() {
+  const { entity } = useEntity();
   const currentSystemName = entity.metadata.name;
   const currentSystemNode = simplifiedEntityName(entity);
   const systemNodes = new Array<{ id: string }>();