From 2f642c7cb2d21cca2f6a7d6ec9a80d5b47ba933a Mon Sep 17 00:00:00 2001
From: Or Ouziel <or.ouziel@elastic.co>
Date: Tue, 31 Jan 2023 22:16:20 +0200
Subject: [PATCH] overload useUIExtension

---
 .../single_page_layout/index.tsx              | 16 ++------
 .../edit_package_policy_page/index.tsx        | 20 +++-------
 .../fleet/public/hooks/use_ui_extension.ts    | 39 ++++++++++++++-----
 3 files changed, 40 insertions(+), 35 deletions(-)

diff --git a/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/create_package_policy_page/single_page_layout/index.tsx b/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/create_package_policy_page/single_page_layout/index.tsx
index c797300cd797e..2ca4c5988b191 100644
--- a/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/create_package_policy_page/single_page_layout/index.tsx
+++ b/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/create_package_policy_page/single_page_layout/index.tsx
@@ -266,18 +266,10 @@ export const CreatePackagePolicySinglePage: CreatePackagePolicyParams = ({
     ]
   );
 
-  const extensionView = useUIExtension(packagePolicy.package?.name ?? '', 'package-policy-create');
-  const replaceDefineStepView = useUIExtension(
-    packagePolicy.package?.name ?? '',
-    'package-policy-replace-define-step'
-  );
-
-  if (replaceDefineStepView && extensionView) {
-    // eslint-disable-next-line no-console
-    console.warn(
-      "'package-policy-create' is ignored when 'package-policy-replace-define-step' is defined"
-    );
-  }
+  const [extensionView, replaceDefineStepView] = useUIExtension(packagePolicy.package?.name ?? '', [
+    'package-policy-create',
+    'package-policy-replace-define-step',
+  ] as const);
 
   const replaceStepConfigurePackagePolicy = replaceDefineStepView && packageInfo?.name && (
     <ExtensionWrapper>
diff --git a/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/edit_package_policy_page/index.tsx b/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/edit_package_policy_page/index.tsx
index 75a69b0cd3160..a76c03b37875e 100644
--- a/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/edit_package_policy_page/index.tsx
+++ b/x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/edit_package_policy_page/index.tsx
@@ -239,23 +239,15 @@ export const EditPackagePolicyForm = memo<{
     }
   };
 
-  const extensionView = useUIExtension(packagePolicy.package?.name ?? '', 'package-policy-edit');
-  const replaceDefineStepView = useUIExtension(
+  const [extensionView, replaceDefineStepView, extensionTabsView] = useUIExtension(
     packagePolicy.package?.name ?? '',
-    'package-policy-replace-define-step'
-  );
-  const extensionTabsView = useUIExtension(
-    packagePolicy.package?.name ?? '',
-    'package-policy-edit-tabs'
+    [
+      'package-policy-edit',
+      'package-policy-replace-define-step',
+      'package-policy-edit-tabs',
+    ] as const
   );
 
-  if (replaceDefineStepView && extensionView) {
-    // eslint-disable-next-line no-console
-    console.warn(
-      "'package-policy-edit' is ignored when 'package-policy-replace-define-step' is defined"
-    );
-  }
-
   const tabsViews = extensionTabsView?.tabs;
   const [selectedTab, setSelectedTab] = useState(0);
 
diff --git a/x-pack/plugins/fleet/public/hooks/use_ui_extension.ts b/x-pack/plugins/fleet/public/hooks/use_ui_extension.ts
index 93e316be241cb..416bf1a72a133 100644
--- a/x-pack/plugins/fleet/public/hooks/use_ui_extension.ts
+++ b/x-pack/plugins/fleet/public/hooks/use_ui_extension.ts
@@ -17,21 +17,42 @@ type NarrowExtensionPoint<V extends UIExtensionPoint['view'], A = UIExtensionPoi
   ? A
   : never;
 
-export const useUIExtension = <V extends UIExtensionPoint['view'] = UIExtensionPoint['view']>(
+type ExclusiveExtensions<T extends ReadonlyArray<UIExtensionPoint['view']>> =
+  | ReplaceWithUndefined<T, 'package-policy-edit' | 'package-policy-create'>
+  | ReplaceWithUndefined<T, 'package-policy-replace-define-step'>;
+
+type ReplaceWithUndefined<Tuple extends readonly unknown[], Id> = Tuple extends readonly []
+  ? []
+  : Tuple extends readonly [infer Head, ...infer Rest]
+  ? Head extends Id
+    ? readonly [undefined, ...ReplaceWithUndefined<Rest, Id>]
+    : Head extends UIExtensionPoint['view']
+    ? readonly [NarrowExtensionPoint<Head> | undefined, ...ReplaceWithUndefined<Rest, Id>]
+    : never
+  : Tuple;
+
+export function useUIExtension<
+  V extends ReadonlyArray<UIExtensionPoint['view']> = ReadonlyArray<UIExtensionPoint['view']>
+>(packageName: UIExtensionPoint['package'], view: V): ExclusiveExtensions<V>;
+
+export function useUIExtension<V extends UIExtensionPoint['view'] = UIExtensionPoint['view']>(
   packageName: UIExtensionPoint['package'],
   view: V
-): NarrowExtensionPoint<V> | undefined => {
-  const registeredExtensions = useContext(UIExtensionsContext);
+): NarrowExtensionPoint<V> | undefined;
 
+export function useUIExtension<V extends UIExtensionPoint['view']>(
+  packageName: UIExtensionPoint['package'],
+  view: V | readonly V[]
+) {
+  const registeredExtensions = useContext(UIExtensionsContext);
   if (!registeredExtensions) {
     throw new Error('useUIExtension called outside of UIExtensionsContext');
   }
 
-  const extension = registeredExtensions?.[packageName]?.[view];
-
-  if (extension) {
-    // FIXME:PT Revisit ignore below and see if TS error can be addressed
-    // @ts-ignore
+  if (typeof view === 'string') {
+    const extension = registeredExtensions?.[packageName]?.[view];
     return extension;
   }
-};
+
+  return view.map((v) => registeredExtensions?.[packageName]?.[v]);
+}