From 888731df92e6264a8f3325b8a4550bd2f76c7d7a Mon Sep 17 00:00:00 2001
From: sjeyaraman <sjeyaraman@expedia.com>
Date: Sun, 5 Jun 2022 12:40:43 -0700
Subject: [PATCH] customize unregister entity menuitem

Signed-off-by: sjeyaraman <sjeyaraman@expedia.com>
---
 .../app/src/components/catalog/EntityPage.tsx | 15 ++++++-
 .../EntityContextMenu.test.tsx                | 25 ++++++++++-
 .../EntityContextMenu/EntityContextMenu.tsx   | 43 +++++++++++++------
 .../components/EntityLayout/EntityLayout.tsx  |  4 +-
 4 files changed, 70 insertions(+), 17 deletions(-)

diff --git a/packages/app/src/components/catalog/EntityPage.tsx b/packages/app/src/components/catalog/EntityPage.tsx
index 1958aabd2e..4318e68e74 100644
--- a/packages/app/src/components/catalog/EntityPage.tsx
+++ b/packages/app/src/components/catalog/EntityPage.tsx
@@ -165,9 +165,22 @@ const EntityLayoutWrapper = (props: { children?: ReactNode }) => {
     ];
   }, []);
 
+  type VisibleType = 'visible' | 'hidden' | 'disabled';
+
+  type contextMenuOptions = {
+    disableUnregister: boolean | VisibleType;
+  };
+
+  const options: contextMenuOptions = {
+    disableUnregister: false,
+  };
+
   return (
     <>
-      <EntityLayout UNSTABLE_extraContextMenuItems={extraMenuItems}>
+      <EntityLayout
+        UNSTABLE_extraContextMenuItems={extraMenuItems}
+        UNSTABLE_contextMenuOptions={options}
+      >
         {props.children}
       </EntityLayout>
       <EntityBadgesDialog
diff --git a/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.test.tsx b/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.test.tsx
index eb920aaa20..9906cd9c12 100644
--- a/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.test.tsx
+++ b/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.test.tsx
@@ -42,7 +42,6 @@ function render(children: React.ReactNode) {
 describe('ComponentContextMenu', () => {
   it('should call onUnregisterEntity on button click', async () => {
     const mockCallback = jest.fn();
-
     await render(
       <EntityContextMenu
         onUnregisterEntity={mockCallback}
@@ -61,6 +60,30 @@ describe('ComponentContextMenu', () => {
     expect(mockCallback).toBeCalled();
   });
 
+  it('check Unregister entity button is disabled', async () => {
+    const mockCallback = jest.fn();
+
+    const { getByText } = await render(
+      <EntityContextMenu
+        UNSTABLE_contextMenuOptions={{ disableUnregister: 'disabled' }}
+        onUnregisterEntity={mockCallback}
+        onInspectEntity={() => {}}
+      />,
+    );
+
+    const button = await screen.findByTestId('menu-button');
+    expect(button).toBeInTheDocument();
+    fireEvent.click(button);
+
+    const unregister = await screen.getByText('Unregister entity');
+    expect(unregister).toBeInTheDocument();
+
+    const unregisterSpanItem = getByText(/Unregister entity/);
+    const unregisterMenuListItem =
+      unregisterSpanItem?.parentElement?.parentElement;
+    expect(unregisterMenuListItem).toHaveAttribute('aria-disabled');
+  });
+
   it('should call onInspectEntity on button click', async () => {
     const mockCallback = jest.fn();
 
diff --git a/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.tsx b/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.tsx
index db50eeab84..0a5e2e1bfe 100644
--- a/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.tsx
+++ b/plugins/catalog/src/components/EntityContextMenu/EntityContextMenu.tsx
@@ -55,9 +55,11 @@ interface ExtraContextMenuItem {
   onClick: () => void;
 }
 
+type VisibleType = 'visible' | 'hidden' | 'disabled';
+
 // unstable context menu option, eg: disable the unregister entity menu
 interface contextMenuOptions {
-  disableUnregister: boolean;
+  disableUnregister: boolean | VisibleType;
 }
 
 interface EntityContextMenuProps {
@@ -106,11 +108,35 @@ export function EntityContextMenu(props: EntityContextMenuProps) {
     <Divider key="the divider is here!" />,
   ];
 
+  const isBoolean =
+    typeof UNSTABLE_contextMenuOptions?.disableUnregister === 'boolean';
+
   const disableUnregister =
     (!unregisterPermission.allowed ||
-      UNSTABLE_contextMenuOptions?.disableUnregister) ??
+      (isBoolean
+        ? UNSTABLE_contextMenuOptions?.disableUnregister
+        : UNSTABLE_contextMenuOptions?.disableUnregister === 'disabled')) ??
     false;
 
+  let unregisterButton = <></>;
+
+  if (UNSTABLE_contextMenuOptions?.disableUnregister !== 'hidden') {
+    unregisterButton = (
+      <MenuItem
+        onClick={() => {
+          onClose();
+          onUnregisterEntity();
+        }}
+        disabled={disableUnregister}
+      >
+        <ListItemIcon>
+          <CancelIcon fontSize="small" />
+        </ListItemIcon>
+        <ListItemText primary="Unregister entity" />
+      </MenuItem>
+    );
+  }
+
   return (
     <>
       <IconButton
@@ -136,18 +162,7 @@ export function EntityContextMenu(props: EntityContextMenuProps) {
       >
         <MenuList>
           {extraItems}
-          <MenuItem
-            onClick={() => {
-              onClose();
-              onUnregisterEntity();
-            }}
-            disabled={disableUnregister}
-          >
-            <ListItemIcon>
-              <CancelIcon fontSize="small" />
-            </ListItemIcon>
-            <ListItemText primary="Unregister entity" />
-          </MenuItem>
+          {unregisterButton}
           <MenuItem
             onClick={() => {
               onClose();
diff --git a/plugins/catalog/src/components/EntityLayout/EntityLayout.tsx b/plugins/catalog/src/components/EntityLayout/EntityLayout.tsx
index 693137fe21..4cbb72db5c 100644
--- a/plugins/catalog/src/components/EntityLayout/EntityLayout.tsx
+++ b/plugins/catalog/src/components/EntityLayout/EntityLayout.tsx
@@ -142,9 +142,11 @@ interface ExtraContextMenuItem {
   onClick: () => void;
 }
 
+type VisibleType = 'visible' | 'hidden' | 'disabled';
+
 // unstable context menu option, eg: disable the unregister entity menu
 interface contextMenuOptions {
-  disableUnregister: boolean;
+  disableUnregister: boolean | VisibleType;
 }
 
 /** @public */