From b0ce104eb1e6b5985c36c5fd90289cfba2dc57e2 Mon Sep 17 00:00:00 2001
From: Aaron <yangtao.yangtao@antgroup.com>
Date: Thu, 22 Aug 2024 20:25:49 +0800
Subject: [PATCH] refactor: getPluginInstance allow match by type secondary
 (#6215)

Co-authored-by: antv <antv@antfin.com>
---
 .../graph/get-plugin-instantce.spec.ts        | 26 +++++++++++++++++--
 packages/g6/src/runtime/plugin.ts             |  9 ++++++-
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/packages/g6/__tests__/unit/runtime/graph/get-plugin-instantce.spec.ts b/packages/g6/__tests__/unit/runtime/graph/get-plugin-instantce.spec.ts
index b3307282080..0b4c0cffbb1 100644
--- a/packages/g6/__tests__/unit/runtime/graph/get-plugin-instantce.spec.ts
+++ b/packages/g6/__tests__/unit/runtime/graph/get-plugin-instantce.spec.ts
@@ -1,4 +1,4 @@
-import { BasePlugin, register } from '@/src';
+import { BasePlugin, ExtensionCategory, register } from '@/src';
 import { createGraph } from '@@/utils';
 
 describe('getPluginInstance', () => {
@@ -12,7 +12,7 @@ describe('getPluginInstance', () => {
       }
     }
 
-    register('plugin', 'custom', CustomPlugin);
+    register(ExtensionCategory.PLUGIN, 'custom', CustomPlugin);
     const graph = createGraph({
       plugins: [
         {
@@ -32,4 +32,26 @@ describe('getPluginInstance', () => {
     const undefinedPlugin = graph.getPluginInstance<CustomPlugin>('undefined-plugin');
     expect(undefinedPlugin).toBe(undefined);
   });
+
+  it('getPluginInstance by type', async () => {
+    const fn = jest.fn();
+
+    class CustomPlugin extends BasePlugin<any> {
+      api() {
+        fn();
+      }
+    }
+
+    register(ExtensionCategory.PLUGIN, 'custom-2', CustomPlugin);
+    const graph = createGraph({
+      plugins: ['custom-2', 'custom-2'],
+    });
+
+    await graph.draw();
+
+    const plugin = graph.getPluginInstance<CustomPlugin>('custom-2');
+    expect(plugin instanceof CustomPlugin).toBe(true);
+    plugin.api();
+    expect(fn).toHaveBeenCalled();
+  });
 });
diff --git a/packages/g6/src/runtime/plugin.ts b/packages/g6/src/runtime/plugin.ts
index b3a405219ae..1a32b29e2d5 100644
--- a/packages/g6/src/runtime/plugin.ts
+++ b/packages/g6/src/runtime/plugin.ts
@@ -1,6 +1,7 @@
 import type { BasePlugin } from '../plugins/base-plugin';
 import { ExtensionController } from '../registry/extension';
 import type { CustomPluginOption, PluginOptions } from '../spec/plugin';
+import { print } from '../utils/print';
 import type { RuntimeContext } from './types';
 
 export class PluginController extends ExtensionController<BasePlugin<CustomPluginOption>> {
@@ -16,6 +17,12 @@ export class PluginController extends ExtensionController<BasePlugin<CustomPlugi
   }
 
   public getPluginInstance(key: string) {
-    return this.extensionMap[key];
+    const exactly = this.extensionMap[key];
+    if (exactly) return exactly;
+
+    print.warn(`Cannot find the plugin ${key}, will try to find it by type.`);
+
+    const fussily = this.extensions.find((extension) => extension.type === key);
+    if (fussily) return this.extensionMap[fussily.key];
   }
 }