From 8230d5910db45c99cd772c27969ec66893073188 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 13 May 2022 16:14:08 +1000
Subject: [PATCH] Rearrange internals of base workspace webviews

---
 extension/src/experiments/workspace.ts   | 17 -----------------
 extension/src/extension.ts               |  2 +-
 extension/src/plots/commands/register.ts | 16 +++++++++++-----
 extension/src/plots/workspace.ts         |  6 ++++++
 extension/src/test/suite/plots/util.ts   |  3 ++-
 extension/src/webview/workspace.ts       | 17 +++++++++++++++++
 6 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/extension/src/experiments/workspace.ts b/extension/src/experiments/workspace.ts
index eb0461ed98..2c9fa08c85 100644
--- a/extension/src/experiments/workspace.ts
+++ b/extension/src/experiments/workspace.ts
@@ -21,8 +21,6 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
 
   public readonly updatesPaused: EventEmitter<boolean>
 
-  private focusedWebviewDvcRoot: string | undefined
-
   constructor(
     internalCommands: InternalCommands,
     updatesPaused: EventEmitter<boolean>,
@@ -40,13 +38,6 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     }
   }
 
-  public getFocusedWebview(): Experiments | undefined {
-    if (!this.focusedWebviewDvcRoot) {
-      return undefined
-    }
-    return this.getRepository(this.focusedWebviewDvcRoot)
-  }
-
   public async addFilter(overrideRoot?: string) {
     const dvcRoot = await this.getDvcRoot(overrideRoot)
     if (!dvcRoot) {
@@ -255,14 +246,6 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return this.runCommand(commandId, cwd, experiment.name)
   }
 
-  private async getDvcRoot(overrideRoot?: string) {
-    return overrideRoot || (await this.getFocusedOrOnlyOrPickProject())
-  }
-
-  private getFocusedOrOnlyOrPickProject() {
-    return this.focusedWebviewDvcRoot || this.getOnlyOrPickProject()
-  }
-
   private pickCurrentExperiment(cwd: string) {
     return this.getRepository(cwd).pickCurrentExperiment()
   }
diff --git a/extension/src/extension.ts b/extension/src/extension.ts
index 4ba8bd6841..ba118bf358 100644
--- a/extension/src/extension.ts
+++ b/extension/src/extension.ts
@@ -221,7 +221,7 @@ export class Extension extends Disposable implements IExtension {
     )
 
     registerExperimentCommands(this.experiments, this.internalCommands)
-    registerPlotsCommands(this.plots)
+    registerPlotsCommands(this.plots, this.internalCommands)
 
     this.dispose.track(
       commands.registerCommand(RegisteredCommands.STOP_EXPERIMENT, async () => {
diff --git a/extension/src/plots/commands/register.ts b/extension/src/plots/commands/register.ts
index 5b73ccfb9f..08c6e9b14a 100644
--- a/extension/src/plots/commands/register.ts
+++ b/extension/src/plots/commands/register.ts
@@ -1,9 +1,15 @@
-import { commands } from 'vscode'
 import { RegisteredCommands } from '../../commands/external'
+import { InternalCommands } from '../../commands/internal'
 import { WorkspacePlots } from '../workspace'
 
-export const registerPlotsCommands = (plots: WorkspacePlots) => {
-  commands.registerCommand(RegisteredCommands.PLOTS_SHOW, () => {
-    plots.showWebview()
-  })
+export const registerPlotsCommands = (
+  plots: WorkspacePlots,
+  internalCommands: InternalCommands
+) => {
+  internalCommands.registerExternalCommand(
+    RegisteredCommands.PLOTS_SHOW,
+    () => {
+      plots.showWebview()
+    }
+  )
 }
diff --git a/extension/src/plots/workspace.ts b/extension/src/plots/workspace.ts
index fb3ed6114e..807dfb177e 100644
--- a/extension/src/plots/workspace.ts
+++ b/extension/src/plots/workspace.ts
@@ -30,6 +30,12 @@ export class WorkspacePlots extends BaseWorkspaceWebviews<Plots, PlotsData> {
       })
     )
 
+    plots.dispose.track(
+      plots.onDidChangeIsWebviewFocused(
+        dvcRoot => (this.focusedWebviewDvcRoot = dvcRoot)
+      )
+    )
+
     return plots
   }
 }
diff --git a/extension/src/test/suite/plots/util.ts b/extension/src/test/suite/plots/util.ts
index 319bd22bc5..b69c6f3c8d 100644
--- a/extension/src/test/suite/plots/util.ts
+++ b/extension/src/test/suite/plots/util.ts
@@ -17,6 +17,7 @@ import { mockHasCheckpoints } from '../experiments/util'
 import { MOCK_IMAGE_MTIME } from '../../fixtures/plotsDiff'
 import { PathsModel } from '../../../plots/paths/model'
 import { Color } from '../../../experiments/model/status/colors'
+import { BaseWorkspaceWebviews } from '../../../webview/workspace'
 
 export const buildPlots = async (
   disposer: Disposer,
@@ -67,7 +68,7 @@ export const buildPlots = async (
 
   await plots.isReady()
 
-  stub(WorkspaceExperiments.prototype, 'getDvcRoots').returns([dvcDemoPath])
+  stub(BaseWorkspaceWebviews.prototype, 'getDvcRoots').returns([dvcDemoPath])
   stub(WorkspaceExperiments.prototype, 'getRepository').returns(experiments)
   stub(WorkspacePlots.prototype, 'getRepository').returns(plots)
 
diff --git a/extension/src/webview/workspace.ts b/extension/src/webview/workspace.ts
index f28f9e66cb..95cfeb3d69 100644
--- a/extension/src/webview/workspace.ts
+++ b/extension/src/webview/workspace.ts
@@ -11,6 +11,8 @@ export abstract class BaseWorkspaceWebviews<
 > extends BaseWorkspace<T, ResourceLocator> {
   protected readonly workspaceState: Memento
 
+  protected focusedWebviewDvcRoot: string | undefined
+
   constructor(
     internalCommands: InternalCommands,
     workspaceState: Memento,
@@ -35,4 +37,19 @@ export abstract class BaseWorkspaceWebviews<
     await repository.showWebview()
     return repository
   }
+
+  public getFocusedWebview(): T | undefined {
+    if (!this.focusedWebviewDvcRoot) {
+      return undefined
+    }
+    return this.getRepository(this.focusedWebviewDvcRoot)
+  }
+
+  protected async getDvcRoot(overrideRoot?: string) {
+    return overrideRoot || (await this.getFocusedOrOnlyOrPickProject())
+  }
+
+  protected getFocusedOrOnlyOrPickProject() {
+    return this.focusedWebviewDvcRoot || this.getOnlyOrPickProject()
+  }
 }