From 83c28783464132804cdfd9848cdf613d1d2a5887 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 9 May 2023 15:40:19 +1000
Subject: [PATCH] add integration test for inline stop action

---
 extension/src/experiments/model/tree.ts       | 46 ++++++++---------
 .../test/suite/experiments/model/tree.test.ts | 49 +++++++++++++++++++
 2 files changed, 72 insertions(+), 23 deletions(-)

diff --git a/extension/src/experiments/model/tree.ts b/extension/src/experiments/model/tree.ts
index 61cb22a521..ebdd1c5c6c 100644
--- a/extension/src/experiments/model/tree.ts
+++ b/extension/src/experiments/model/tree.ts
@@ -124,30 +124,10 @@ export class ExperimentsTree
   }
 
   private registerWorkaroundCommands() {
-    const callCommandWithSelected = async (
-      command:
-        | RegisteredCliCommands.EXPERIMENT_VIEW_REMOVE
-        | RegisteredCliCommands.EXPERIMENT_VIEW_PUSH
-        | RegisteredCommands.EXPERIMENT_VIEW_STOP,
-      experimentItem: ExperimentItem | string,
-      types: ExperimentType[]
-    ) => {
-      const selected = [
-        ...this.getSelectedExperimentItems(),
-        experimentItem
-      ] as (string | ExperimentItem)[]
-
-      const acc = collectExperimentType(selected, new Set(types))
-
-      for (const [dvcRoot, ids] of Object.entries(acc)) {
-        await commands.executeCommand(command, { dvcRoot, ids: [...ids] })
-      }
-    }
-
     commands.registerCommand(
       'dvc.views.experimentsTree.removeExperiment',
       (experimentItem: ExperimentItem) =>
-        callCommandWithSelected(
+        this.callCommandWithSelected(
           RegisteredCliCommands.EXPERIMENT_VIEW_REMOVE,
           experimentItem,
           [ExperimentType.EXPERIMENT, ExperimentType.QUEUED]
@@ -157,7 +137,7 @@ export class ExperimentsTree
     commands.registerCommand(
       'dvc.views.experimentsTree.pushExperiment',
       (experimentItem: ExperimentItem) =>
-        callCommandWithSelected(
+        this.callCommandWithSelected(
           RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
           experimentItem,
           [ExperimentType.EXPERIMENT]
@@ -167,7 +147,7 @@ export class ExperimentsTree
     commands.registerCommand(
       'dvc.views.experimentsTree.stopExperiment',
       (experimentItem: ExperimentItem) =>
-        callCommandWithSelected(
+        this.callCommandWithSelected(
           RegisteredCommands.EXPERIMENT_VIEW_STOP,
           experimentItem,
           [ExperimentType.RUNNING]
@@ -175,6 +155,26 @@ export class ExperimentsTree
     )
   }
 
+  private async callCommandWithSelected(
+    command:
+      | RegisteredCliCommands.EXPERIMENT_VIEW_REMOVE
+      | RegisteredCliCommands.EXPERIMENT_VIEW_PUSH
+      | RegisteredCommands.EXPERIMENT_VIEW_STOP,
+    experimentItem: ExperimentItem | string,
+    types: ExperimentType[]
+  ) {
+    const selected = [...this.getSelectedExperimentItems(), experimentItem] as (
+      | string
+      | ExperimentItem
+    )[]
+
+    const acc = collectExperimentType(selected, new Set(types))
+
+    for (const [dvcRoot, ids] of Object.entries(acc)) {
+      await commands.executeCommand(command, { dvcRoot, ids: [...ids] })
+    }
+  }
+
   private async getRootElements() {
     await this.experiments.isReady()
     const dvcRoots = this.experiments.getDvcRoots()
diff --git a/extension/src/test/suite/experiments/model/tree.test.ts b/extension/src/test/suite/experiments/model/tree.test.ts
index 9b3ad3764b..5d2cc1833d 100644
--- a/extension/src/test/suite/experiments/model/tree.test.ts
+++ b/extension/src/test/suite/experiments/model/tree.test.ts
@@ -439,6 +439,55 @@ suite('Experiments Tree Test Suite', () => {
       )
     })
 
+    it('should be able to stop multiple running experiments with dvc.views.experimentsTree.stopExperiment', async () => {
+      bypassProgressCloseDelay()
+      const mockFirstExperimentId = 'first-exp-stopped'
+      const mockSecondExperimentId = 'second-exp-stopped'
+      const mockQueuedExperimentLabel = 'queued-excluded'
+
+      const mockStopExperiments = stub(
+        WorkspaceExperiments.prototype,
+        'stopExperiments'
+      ).resolves(undefined)
+
+      stubPrivatePrototypeMethod(
+        ExperimentsTree,
+        'getSelectedExperimentItems'
+      ).returns([
+        dvcDemoPath,
+        {
+          dvcRoot: dvcDemoPath,
+          label: mockQueuedExperimentLabel,
+          type: ExperimentType.QUEUED
+        },
+        {
+          dvcRoot: dvcDemoPath,
+          id: mockFirstExperimentId,
+          type: ExperimentType.RUNNING
+        },
+        {
+          dvcRoot: dvcDemoPath,
+          id: 'workspace-excluded',
+          type: ExperimentType.WORKSPACE
+        }
+      ])
+
+      await commands.executeCommand(
+        'dvc.views.experimentsTree.stopExperiment',
+        {
+          dvcRoot: dvcDemoPath,
+          id: mockSecondExperimentId,
+          type: ExperimentType.RUNNING
+        }
+      )
+
+      expect(mockStopExperiments).to.be.calledWithExactly(
+        dvcDemoPath,
+        mockFirstExperimentId,
+        mockSecondExperimentId
+      )
+    })
+
     it('should be able to apply an experiment to the workspace with dvc.views.experiments.applyExperiment', async () => {
       const { experiments } = buildExperiments(disposable)