From 704d190c85a1339c54b23d57ab2cd858275f65c3 Mon Sep 17 00:00:00 2001
From: Shailendra Ahir <shailendraahir@hotmail.com>
Date: Wed, 23 Mar 2022 08:50:05 +0530
Subject: [PATCH] #8850 Added prop types and test cases

Signed-off-by: Shailendra Ahir <shailendraahir@hotmail.com>
---
 .../EntityRelationsGraph.test.tsx             | 47 +++++++++++++++++++
 .../EntityRelationsGraph.tsx                  |  6 +--
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.test.tsx b/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.test.tsx
index 2bf303bb92..5d9e9b18a6 100644
--- a/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.test.tsx
+++ b/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.test.tsx
@@ -21,6 +21,7 @@ import {
   RELATION_OWNER_OF,
   RELATION_PART_OF,
 } from '@backstage/catalog-model';
+import { DependencyGraphTypes } from '@backstage/core-components';
 import { CatalogApi, catalogApiRef } from '@backstage/plugin-catalog-react';
 import { renderInTestApp, TestApiProvider } from '@backstage/test-utils';
 import userEvent from '@testing-library/user-event';
@@ -30,6 +31,7 @@ import { EntityRelationsGraph } from './EntityRelationsGraph';
 describe('<EntityRelationsGraph/>', () => {
   let Wrapper: FunctionComponent;
   let catalog: jest.Mocked<CatalogApi>;
+  const CUSTOM_TEST_ID = 'custom-test-id';
 
   beforeAll(() => {
     Object.defineProperty(window.SVGElement.prototype, 'getBBox', {
@@ -378,4 +380,49 @@ describe('<EntityRelationsGraph/>', () => {
     userEvent.click(await findByText('k:d/a1'));
     expect(onNodeClick).toBeCalledTimes(1);
   });
+
+  test('render custom node', async () => {
+    const renderNode = (props: DependencyGraphTypes.RenderNodeProps) => (
+      <g>
+        <text>{props.node.id}</text>
+        <circle data-testid={CUSTOM_TEST_ID} r={100} />
+      </g>
+    );
+
+    const { findAllByTestId, container } = await renderInTestApp(
+      <Wrapper>
+        <EntityRelationsGraph
+          rootEntityNames={{ kind: 'b', namespace: 'd', name: 'c' }}
+          renderNode={renderNode}
+        />
+      </Wrapper>,
+    );
+
+    const node = await findAllByTestId(CUSTOM_TEST_ID);
+    expect(node[0]).toBeInTheDocument();
+    expect(container.querySelector('circle')).toBeInTheDocument();
+  });
+
+  test('render custom label', async () => {
+    const renderLabel = (props: DependencyGraphTypes.RenderLabelProps) => (
+      <g>
+        <text>{`Test-Label${props.edge.label}`}</text>
+        <circle data-testid={CUSTOM_TEST_ID} r={100} />
+      </g>
+    );
+
+    const { findAllByTestId, findAllByText, container } = await renderInTestApp(
+      <Wrapper>
+        <EntityRelationsGraph
+          rootEntityNames={{ kind: 'b', namespace: 'd', name: 'c' }}
+          renderLabel={renderLabel}
+        />
+      </Wrapper>,
+    );
+    const node = await findAllByTestId(CUSTOM_TEST_ID);
+    expect(node[0]).toBeInTheDocument();
+    expect(container.querySelector('circle')).toBeInTheDocument();
+    const labels = await findAllByText('Test-Labelvisible');
+    expect(labels[0]).toBeInTheDocument();
+  });
 });
diff --git a/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.tsx b/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.tsx
index 19ca5398a6..bf3fbf8f3a 100644
--- a/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.tsx
+++ b/plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.tsx
@@ -28,7 +28,7 @@ import React, { MouseEvent, useEffect, useMemo } from 'react';
 import { CustomLabel } from './CustomLabel';
 import { CustomNode } from './CustomNode';
 import { ALL_RELATION_PAIRS, RelationPairs } from './relations';
-import { Direction, EntityNode } from './types';
+import { Direction, EntityEdge, EntityNode } from './types';
 import { useEntityRelationNodesAndEdges } from './useEntityRelationNodesAndEdges';
 
 const useStyles = makeStyles(theme => ({
@@ -93,8 +93,8 @@ export const EntityRelationsGraph = ({
   relationPairs?: RelationPairs;
   className?: string;
   zoom?: 'enabled' | 'disabled' | 'enable-on-click';
-  renderNode?: any;
-  renderLabel?: any;
+  renderNode?: DependencyGraphTypes.RenderNodeFunction<EntityNode>;
+  renderLabel?: DependencyGraphTypes.RenderLabelFunction<EntityEdge>;
 }) => {
   const theme = useTheme();
   const classes = useStyles();