Add behaviour for tray menu items originating from extensions (#5609)

src/renderer/components/test-utils/get-application-builder.tsx

@@ -6,7 +6,7 @@ import type { LensRendererExtension } from "../../../extensions/lens-renderer-ex
 import rendererExtensionsInjectable from "../../../extensions/renderer-extensions.injectable";
 import currentlyInClusterFrameInjectable from "../../routes/currently-in-cluster-frame.injectable";
 import { extensionRegistratorInjectionToken } from "../../../extensions/extension-loader/extension-registrator-injection-token";
-import type { IObservableArray } from "mobx";
+import type { IObservableArray, ObservableSet } from "mobx";
 import { computed, observable, runInAction } from "mobx";
 import { renderFor } from "./renderFor";
 import React from "react";
@@ -24,7 +24,7 @@ import type { ClusterStore } from "../../../common/cluster-store/cluster-store";
 import mainExtensionsInjectable from "../../../extensions/main-extensions.injectable";
 import currentRouteComponentInjectable from "../../routes/current-route-component.injectable";
 import { pipeline } from "@ogre-tools/fp";
-import { flatMap, compact, join, get, filter, map, matches, find } from "lodash/fp";
+import { flatMap, compact, join, get, filter, map, matches } from "lodash/fp";
 import preferenceNavigationItemsInjectable from "../+preferences/preferences-navigation/preference-navigation-items.injectable";
 import navigateToPreferencesInjectable from "../../../common/front-end-routing/routes/preferences/navigate-to-preferences.injectable";
 import type { MenuItemOpts } from "../../../main/menu/application-menu-items.injectable";
@@ -56,13 +56,31 @@ import { openMenu } from "react-select-event";
 import userEvent from "@testing-library/user-event";
 import { StatusBar } from "../status-bar/status-bar";
 import lensProxyPortInjectable from "../../../main/lens-proxy/lens-proxy-port.injectable";
+import type { LensMainExtension } from "../../../extensions/lens-main-extension";
+import trayMenuItemsInjectable from "../../../main/tray/tray-menu-item/tray-menu-items.injectable";
+import type { LensExtension } from "../../../extensions/lens-extension";
 
 type Callback = (dis: DiContainers) => void | Promise<void>;
 
+type EnableExtensions<T> = (...extensions: T[]) => Promise<void>;
+type DisableExtensions<T> = (...extensions: T[]) => void;
+
 export interface ApplicationBuilder {
   dis: DiContainers;
   setEnvironmentToClusterFrame: () => ApplicationBuilder;
-  addExtensions: (...extensions: LensRendererExtension[]) => Promise<ApplicationBuilder>;
+
+  extensions: {
+    renderer: {
+      enable: EnableExtensions<LensRendererExtension>;
+      disable: DisableExtensions<LensRendererExtension>;
+    };
+
+    main: {
+      enable: EnableExtensions<LensMainExtension>;
+      disable: DisableExtensions<LensMainExtension>;
+    };
+  };
+
   allowKubeResource: (resourceName: KubeResource) => ApplicationBuilder;
   beforeApplicationStart: (callback: Callback) => ApplicationBuilder;
   beforeRender: (callback: Callback) => ApplicationBuilder;
@@ -135,7 +153,8 @@ export const getApplicationBuilder = () => {
   const beforeApplicationStartCallbacks: Callback[] = [];
   const beforeRenderCallbacks: Callback[] = [];
 
-  const extensionsState = observable.array<LensRendererExtension>();
+  const rendererExtensionsState = observable.set<LensRendererExtension>();
+  const mainExtensionsState = observable.set<LensMainExtension>();
 
   rendererDi.override(subscribeStoresInjectable, () => () => () => {});
 
@@ -174,14 +193,13 @@ export const getApplicationBuilder = () => {
   );
 
   rendererDi.override(rendererExtensionsInjectable, () =>
-    computed(() => extensionsState),
+    computed(() => [...rendererExtensionsState]),
   );
 
   mainDi.override(mainExtensionsInjectable, () =>
-    computed(() => []),
+    computed(() => [...mainExtensionsState]),
   );
 
-  let trayMenuItemsStateFake: TrayMenuItem[];
   let trayMenuIconPath: string;
 
   mainDi.override(electronTrayInjectable, () => ({
@@ -191,9 +209,7 @@ export const getApplicationBuilder = () => {
       trayMenuIconPath = iconPaths.normal;
     },
     stop: () => {},
-    setMenuItems: (items) => {
-      trayMenuItemsStateFake = items;
-    },
+    setMenuItems: () => {},
     setIconPath: (path) => {
       trayMenuIconPath = path;
     },
@@ -202,6 +218,43 @@ export const getApplicationBuilder = () => {
   let allowedResourcesState: IObservableArray<KubeResource>;
   let rendered: RenderResult;
 
+  const enableExtensionsFor = <T extends ObservableSet>(
+    extensionState: T,
+    di: DiContainer,
+  ) => {
+    let index = 0;
+
+    return async (...extensions: LensExtension[]) => {
+      const extensionRegistrators = di.injectMany(
+        extensionRegistratorInjectionToken,
+      );
+
+      const addAndEnableExtensions = async () => {
+        index++;
+
+        const registratorPromises = extensions.flatMap((extension) =>
+          extensionRegistrators.map((registrator) =>
+            registrator(extension, index),
+          ),
+        );
+
+        await Promise.all(registratorPromises);
+
+        runInAction(() => {
+          extensions.forEach((extension) => {
+            extensionState.add(extension);
+          });
+        });
+      };
+
+      if (rendered) {
+        await addAndEnableExtensions();
+      } else {
+        builder.beforeRender(addAndEnableExtensions);
+      }
+    };
+  };
+
   const builder: ApplicationBuilder = {
     dis,
 
@@ -242,18 +295,20 @@ export const getApplicationBuilder = () => {
 
     tray: {
       get: (id: string) => {
-        return trayMenuItemsStateFake.find(matches({ id })) ?? null;
+        const trayMenuItems = mainDi.inject(trayMenuItemsInjectable).get();
+
+        return trayMenuItems.find(matches({ id })) ?? null;
       },
       getIconPath: () => trayMenuIconPath,
+
       click: async (id: string) => {
-        const menuItem = pipeline(
-          trayMenuItemsStateFake,
-          find((menuItem) => menuItem.id === id),
-        );
+        const trayMenuItems = mainDi.inject(trayMenuItemsInjectable).get();
+
+        const menuItem = trayMenuItems.find(matches({ id })) ?? null;
 
         if (!menuItem) {
           const availableIds = pipeline(
-            trayMenuItemsStateFake,
+            trayMenuItems,
             filter(item => !!item.click),
             map(item => item.id),
             join(", "),
@@ -345,32 +400,16 @@ export const getApplicationBuilder = () => {
       return builder;
     },
 
-    addExtensions: async (...extensions) => {
-      const extensionRegistrators = rendererDi.injectMany(
-        extensionRegistratorInjectionToken,
-      );
+    extensions: {
+      renderer: {
+        enable: enableExtensionsFor(rendererExtensionsState, rendererDi),
+        disable: disableExtensionsFor(rendererExtensionsState),
+      },
 
-      const addAndEnableExtensions = async () => {
-        const registratorPromises = extensions.flatMap((extension) =>
-          extensionRegistrators.map((registrator) => registrator(extension, 1)),
-        );
-
-        await Promise.all(registratorPromises);
-
-        runInAction(() => {
-          extensions.forEach((extension) => {
-            extensionsState.push(extension);
-          });
-        });
-      };
-
-      if (rendered) {
-        await addAndEnableExtensions();
-      } else {
-        builder.beforeRender(addAndEnableExtensions);
-      }
-
-      return builder;
+      main: {
+        enable: enableExtensionsFor(mainExtensionsState, mainDi),
+        disable: disableExtensionsFor(mainExtensionsState),
+      },
     },
 
     allowKubeResource: (resourceName) => {
@@ -494,3 +533,14 @@ function toFlatChildren(parentId: string | null | undefined): ToFlatChildren {
     ),
   ];
 }
+
+const disableExtensionsFor =
+    <T extends ObservableSet>(extensionState: T) =>
+
+    (...extensions: LensExtension[]) => {
+      extensions.forEach((extension) => {
+        runInAction(() => {
+          extensionState.delete(extension);
+        });
+      });
+    };