From f3c4b71c8dc139604f22838d057869ca18459cfa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torkel=20=C3=96degaard?= <torkel@grafana.com>
Date: Fri, 13 Sep 2024 08:50:13 +0200
Subject: [PATCH] VizPanel: Harden plugin change handling

---
 .../src/components/VizPanel/VizPanel.test.tsx | 20 ++++++++++++-
 .../src/components/VizPanel/VizPanel.tsx      | 30 ++++++++++++-------
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/packages/scenes/src/components/VizPanel/VizPanel.test.tsx b/packages/scenes/src/components/VizPanel/VizPanel.test.tsx
index be1717e3d..fb2dc9914 100644
--- a/packages/scenes/src/components/VizPanel/VizPanel.test.tsx
+++ b/packages/scenes/src/components/VizPanel/VizPanel.test.tsx
@@ -24,6 +24,7 @@ import { SeriesVisibilityChangeMode } from '@grafana/ui';
 import { SceneTimeRange } from '../../core/SceneTimeRange';
 import { act, render, screen } from '@testing-library/react';
 import { RefreshEvent } from '@grafana/runtime';
+import { SceneDeactivationHandler } from '../../core/types';
 
 let pluginToLoad: PanelPlugin | undefined;
 
@@ -289,6 +290,7 @@ describe('VizPanel', () => {
 
   describe('When changing plugin', () => {
     let panel: VizPanel<OptionsPlugin1, FieldConfigPlugin1>;
+    let deactivate: SceneDeactivationHandler;
 
     beforeEach(async () => {
       panel = new VizPanel<OptionsPlugin1, FieldConfigPlugin1>({
@@ -300,7 +302,11 @@ describe('VizPanel', () => {
       });
 
       pluginToLoad = getTestPlugin1();
-      panel.activate();
+      deactivate = panel.activate();
+    });
+
+    afterEach(() => {
+      deactivate();
     });
 
     it('Should successfully change from one viz type to another', async () => {
@@ -365,6 +371,18 @@ describe('VizPanel', () => {
       expect(panel.state.options.showThresholds).toBe(true);
       expect(panel.state.options.option2).toBe('hello');
     });
+
+    it('Should detect plugin change on activation', async () => {
+      pluginToLoad = getTestPlugin2();
+      deactivate();
+
+      panel.setState({ pluginId: pluginToLoad.meta.id });
+      deactivate = panel.activate();
+
+      await Promise.resolve();
+
+      expect(panel.getPlugin()).toBe(pluginToLoad);
+    });
   });
 
   describe('getLegacyPanelId', () => {
diff --git a/packages/scenes/src/components/VizPanel/VizPanel.tsx b/packages/scenes/src/components/VizPanel/VizPanel.tsx
index b02664793..61c1e9c1e 100644
--- a/packages/scenes/src/components/VizPanel/VizPanel.tsx
+++ b/packages/scenes/src/components/VizPanel/VizPanel.tsx
@@ -120,12 +120,17 @@ export class VizPanel<TOptions = {}, TFieldConfig extends {} = {}> extends Scene
   }
 
   private _onActivate() {
-    if (!this._plugin) {
+    if (!this._plugin || this.state.pluginId !== this._plugin.meta.id) {
       this._loadPlugin(this.state.pluginId);
     }
   }
 
-  private async _loadPlugin(pluginId: string, overwriteOptions?: DeepPartial<{}>, overwriteFieldConfig?: FieldConfigSource, isAfterPluginChange?: boolean) {
+  private async _loadPlugin(
+    pluginId: string,
+    overwriteOptions?: DeepPartial<{}>,
+    overwriteFieldConfig?: FieldConfigSource,
+    isAfterPluginChange?: boolean
+  ) {
     const plugin = loadPanelPluginSync(pluginId);
 
     if (plugin) {
@@ -155,7 +160,12 @@ export class VizPanel<TOptions = {}, TFieldConfig extends {} = {}> extends Scene
     return panelId;
   }
 
-  private async _pluginLoaded(plugin: PanelPlugin, overwriteOptions?: DeepPartial<{}>, overwriteFieldConfig?: FieldConfigSource, isAfterPluginChange?: boolean) {
+  private async _pluginLoaded(
+    plugin: PanelPlugin,
+    overwriteOptions?: DeepPartial<{}>,
+    overwriteFieldConfig?: FieldConfigSource,
+    isAfterPluginChange?: boolean
+  ) {
     const { options, fieldConfig, title, pluginVersion, _UNSAFE_customMigrationHandler } = this.state;
 
     const panel: PanelModel = {
@@ -215,6 +225,10 @@ export class VizPanel<TOptions = {}, TFieldConfig extends {} = {}> extends Scene
     return this._plugin;
   }
 
+  public getPluginAsync(): PanelPlugin | undefined {
+    return this._plugin;
+  }
+
   public getPanelContext(): PanelContext {
     this._panelContext ??= this.buildPanelContext();
 
@@ -255,11 +269,7 @@ export class VizPanel<TOptions = {}, TFieldConfig extends {} = {}> extends Scene
   };
 
   public async changePluginType(pluginId: string, newOptions?: DeepPartial<{}>, newFieldConfig?: FieldConfigSource) {
-    const {
-      options: prevOptions,
-      fieldConfig: prevFieldConfig,
-      pluginId: prevPluginId,
-    } = this.state;
+    const { options: prevOptions, fieldConfig: prevFieldConfig, pluginId: prevPluginId } = this.state;
 
     //clear field config cache to update it later
     this._dataWithFieldConfig = undefined;
@@ -274,8 +284,8 @@ export class VizPanel<TOptions = {}, TFieldConfig extends {} = {}> extends Scene
       type: pluginId,
     };
 
-    // onPanelTypeChanged is mainly used by plugins to migrate from Angular to React. 
-    // For example, this will migrate options from 'graph' to 'timeseries' if the previous and new plugin ID matches. 
+    // onPanelTypeChanged is mainly used by plugins to migrate from Angular to React.
+    // For example, this will migrate options from 'graph' to 'timeseries' if the previous and new plugin ID matches.
     const updatedOptions = this._plugin?.onPanelTypeChanged?.(panel, prevPluginId, prevOptions, prevFieldConfig);
 
     if (updatedOptions && !isEmpty(updatedOptions)) {