Refactor stubWorkspaceExperimentsGetters

iterative · Jul 14, 2023 · 664987e · 664987e
1 parent 207a11f
commit 664987e
Show file tree

Hide file tree

Showing 6 changed files with 87 additions and 164 deletions.
diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts
@@ -533,7 +533,7 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to apply an experiment', async () => {
-      const { experiments } = buildExperiments({ disposer: disposable })
+      const { experiments } = stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
 
       const webview = await experiments.showWebview()
@@ -545,8 +545,6 @@ suite('Experiments Test Suite', () => {
         'expApply'
       ).resolves(undefined)
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       mockMessageReceived.fire({
         payload: mockExperimentId,
         type: MessageFromWebviewType.APPLY_EXPERIMENT_TO_WORKSPACE
@@ -560,7 +558,7 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to create a branch from an experiment', async () => {
-      const { experiments } = buildExperiments({ disposer: disposable })
+      const { experiments } = stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
 
       const mockBranch = 'mock-branch-input'
@@ -571,8 +569,6 @@ suite('Experiments Test Suite', () => {
         'expBranch'
       ).resolves('undefined')
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const webview = await experiments.showWebview()
       const mockMessageReceived = getMessageReceivedEmitter(webview)
       const mockExperimentId = 'exp-e7a67'
@@ -714,9 +710,8 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to modify the workspace params and queue an experiment', async () => {
-      const { experiments, dvcExecutor } = buildExperiments({
-        disposer: disposable
-      })
+      const { experiments, dvcExecutor } =
+        stubWorkspaceExperimentsGetters(disposable)
 
       const mockModifiedParams = [
         '-S',
@@ -725,8 +720,6 @@ suite('Experiments Test Suite', () => {
         'params.yaml:weight_decay=0'
       ]
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       stub(experiments, 'pickAndModifyParams').resolves(mockModifiedParams)
       const mockQueueExperiment = stub(dvcExecutor, 'expRunQueue').resolves(
         undefined
@@ -749,9 +742,8 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to modify the workspace params and run a new experiment', async () => {
-      const { experiments, dvcRunner } = buildExperiments({
-        disposer: disposable
-      })
+      const { experiments, dvcRunner } =
+        stubWorkspaceExperimentsGetters(disposable)
 
       const mockModifiedParams = [
         '-S',
@@ -760,8 +752,6 @@ suite('Experiments Test Suite', () => {
         'params.yaml:weight_decay=0'
       ]
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       stub(experiments, 'pickAndModifyParams').resolves(mockModifiedParams)
       const mockRunExperiment = stub(dvcRunner, 'runExperiment').resolves(
         undefined
@@ -786,9 +776,8 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to modify the workspace params, reset and run a new experiment', async () => {
-      const { experiments, dvcRunner } = buildExperiments({
-        disposer: disposable
-      })
+      const { experiments, dvcRunner } =
+        stubWorkspaceExperimentsGetters(disposable)
 
       const mockModifiedParams = [
         '-S',
@@ -799,8 +788,6 @@ suite('Experiments Test Suite', () => {
 
       stub(experiments, 'pickAndModifyParams').resolves(mockModifiedParams)
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const webview = await experiments.showWebview()
       const mockMessageReceived = getMessageReceivedEmitter(webview)
       const mockRunExperiment = stub(dvcRunner, 'runExperiment').resolves(
@@ -847,9 +834,8 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it("should be able to handle a message to toggle an experiment's status", async () => {
-      const { experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
+      const { experiments, experimentsModel } =
+        stubWorkspaceExperimentsGetters(disposable)
 
       await experiments.isReady()
 
@@ -875,8 +861,6 @@ suite('Experiments Test Suite', () => {
         'queued experiment cannot be selected'
       ).to.be.false
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const webview = await experiments.showWebview()
       const mockMessageReceived = getMessageReceivedEmitter(webview)
       const toggleSpy = spy(experimentsModel, 'toggleStatus')
@@ -1192,9 +1176,8 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to select experiments for plotting', async () => {
-      const { experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
+      const { experiments, experimentsModel } =
+        stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
 
       const webview = await experiments.showWebview()
@@ -1204,8 +1187,6 @@ suite('Experiments Test Suite', () => {
       const expectedIds = ['main', 'test-branch', runningInQueueId]
       const mockExperimentIds = [...expectedIds, queuedId]
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const tableChangePromise = experimentsUpdatedEvent(experiments)
 
       mockMessageReceived.fire({
@@ -1231,9 +1212,8 @@ suite('Experiments Test Suite', () => {
         dvc: true,
         experiments: true
       })
-      const { experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
+      const { experiments, experimentsModel } =
+        stubWorkspaceExperimentsGetters(disposable)
       const mockShowPlots = stub(WorkspacePlots.prototype, 'showWebview')
 
       const dataSent = new Promise(resolve =>
@@ -1250,8 +1230,6 @@ suite('Experiments Test Suite', () => {
       const runningInQueueId = 'exp-e7a67'
       const mockExperimentIds = ['main', 'test-branch', runningInQueueId]
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const tableChangePromise = experimentsUpdatedEvent(experiments)
 
       mockMessageReceived.fire({
@@ -1302,9 +1280,9 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should handle a message to stop experiments running', async () => {
-      const { experiments, dvcExecutor } = buildExperiments({
-        disposer: disposable
-      })
+      const { experiments, dvcExecutor } =
+        stubWorkspaceExperimentsGetters(disposable)
+
       const mockQueueKill = stub(dvcExecutor, 'queueKill')
       const mockStopProcesses = stub(ProcessExecution, 'stopProcesses')
 
@@ -1328,7 +1306,6 @@ suite('Experiments Test Suite', () => {
       const mockMessageReceived = getMessageReceivedEmitter(webview)
       const mockExperimentIds = ['exp-e7a67', 'exp-83425']
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
       const mockPid = 1234
       const mockGetPidFromFile = stub(FileSystem, 'getPidFromFile')
         .onFirstCall()

diff --git a/extension/src/test/suite/experiments/model/filterBy/tree.test.ts b/extension/src/test/suite/experiments/model/filterBy/tree.test.ts
@@ -60,15 +60,11 @@ suite('Experiments Filter By Tree Test Suite', () => {
     })
 
     it('should be able to update the table data by adding and removing a filter', async () => {
-      const { experiments, messageSpy } = buildExperiments({
-        disposer: disposable
-      })
-
+      const { experiments, messageSpy } =
+        stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
       await experiments.showWebview()
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const accuracyPath = buildMetricOrParamPath(
         ColumnType.METRICS,
         'summary.json',
@@ -135,14 +131,11 @@ suite('Experiments Filter By Tree Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to remove all filters with dvc.views.experimentsFilterByTree.removeAllFilters', async () => {
-      const mockShowQuickPick = stub(window, 'showQuickPick')
-      const mockShowInputBox = stub(window, 'showInputBox')
-
-      const { experiments } = buildExperiments({ disposer: disposable })
-
+      const { experiments } = stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
+      const mockShowQuickPick = stub(window, 'showQuickPick')
+      const mockShowInputBox = stub(window, 'showInputBox')
 
       const lossPath = buildMetricOrParamPath(
         ColumnType.METRICS,
@@ -338,15 +331,11 @@ suite('Experiments Filter By Tree Test Suite', () => {
     })
 
     it('should be able to filter to starred experiments', async () => {
-      const { experiments, messageSpy } = buildExperiments({
-        disposer: disposable
-      })
-
+      const { experiments, messageSpy } =
+        stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
       await experiments.showWebview()
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       await addFilterViaQuickInput(experiments, starredFilter)
 
       const [workspace, main, fe2919b, _7df876c] = rowsFixture
@@ -373,14 +362,10 @@ suite('Experiments Filter By Tree Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should provide a shortcut to filter to starred experiments', async () => {
-      const { experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
-
+      const { experiments, experimentsModel } =
+        stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const mockAddFilter = stub(experimentsModel, 'addFilter')
 
       await commands.executeCommand(

diff --git a/extension/src/test/suite/experiments/model/sortBy/tree.test.ts b/extension/src/test/suite/experiments/model/sortBy/tree.test.ts
@@ -164,7 +164,10 @@ suite('Experiments Sort By Tree Test Suite', () => {
           )
 
       stub(WorkspaceExperiments.prototype, 'getDvcRoots').returns([dvcDemoPath])
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
+      stub(WorkspaceExperiments.prototype, 'getOnlyOrPickProject').resolves(
+        dvcDemoPath
+      )
+      stub(WorkspaceExperiments.prototype, 'getRepository').returns(experiments)
 
       mockSortQuickPicks(testParamPath, false)
       const tableChangedPromise = experimentsUpdatedEvent(experiments)
@@ -284,14 +287,10 @@ suite('Experiments Sort By Tree Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should provide a shortcut to sort by starred experiments', async () => {
-      const { experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
-
+      const { experiments, experimentsModel } =
+        stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
 
-      stubWorkspaceExperimentsGetters(dvcDemoPath, experiments)
-
       const mockAddSort = stub(experimentsModel, 'addSort')
 
       await commands.executeCommand(

diff --git a/extension/src/test/suite/experiments/model/tree.test.ts b/extension/src/test/suite/experiments/model/tree.test.ts
@@ -573,21 +573,19 @@ suite('Experiments Tree Test Suite', () => {
     })
 
     it('should be able to queue an experiment from an existing one with dvc.views.experiments.queueExperiment', async () => {
-      const { dvcExecutor, experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
+      const {
+        dvcExecutor,
+        experiments,
+        experimentsModel,
+        mockGetOnlyOrPickProject
+      } = stubWorkspaceExperimentsGetters(disposable)
 
       await experiments.isReady()
 
       const mockExperimentRunQueue = stub(dvcExecutor, 'expRunQueue').resolves(
         'true'
       )
 
-      const [mockGetOnlyOrPickProject] = stubWorkspaceExperimentsGetters(
-        '',
-        experiments
-      )
-
       const getParamsSpy = spy(experimentsModel, 'getWorkspaceParams')
 
       const mockShowQuickPick = stub(window, 'showQuickPick') as SinonStub<
@@ -632,21 +630,19 @@ suite('Experiments Tree Test Suite', () => {
     })
 
     it('should be able to run a new experiment from an existing one with dvc.views.experiments.runExperiment', async () => {
-      const { dvcRunner, experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
+      const {
+        dvcRunner,
+        experiments,
+        experimentsModel,
+        mockGetOnlyOrPickProject
+      } = stubWorkspaceExperimentsGetters(disposable)
 
       await experiments.isReady()
 
       const mockRunExperiment = stub(dvcRunner, 'runExperiment').resolves(
         undefined
       )
 
-      const [mockGetOnlyOrPickProject] = stubWorkspaceExperimentsGetters(
-        '',
-        experiments
-      )
-
       const getParamsSpy = spy(experimentsModel, 'getWorkspaceParams')
 
       const mockShowQuickPick = stub(window, 'showQuickPick') as SinonStub<
@@ -688,22 +684,19 @@ suite('Experiments Tree Test Suite', () => {
     })
 
     it('should be able to reset and run a new checkpoint experiment from an existing one with dvc.views.experiments.resetAndRunCheckpointExperiment', async () => {
-      const { dvcRunner, experiments, experimentsModel } = buildExperiments({
-        disposer: disposable
-      })
-
+      const {
+        dvcRunner,
+        experiments,
+        experimentsModel,
+        mockGetOnlyOrPickProject
+      } = stubWorkspaceExperimentsGetters(disposable)
       await experiments.isReady()
 
       const mockRunExperimentReset = stub(
         dvcRunner,
         'runExperimentReset'
       ).resolves(undefined)
 
-      const [mockGetOnlyOrPickProject] = stubWorkspaceExperimentsGetters(
-        '',
-        experiments
-      )
-
       const getParamsSpy = spy(experimentsModel, 'getWorkspaceParams')
 
       const mockShowQuickPick = stub(window, 'showQuickPick') as SinonStub<

diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts
@@ -244,9 +244,12 @@ export const buildExperimentsData = (
 }
 
 export const stubWorkspaceExperimentsGetters = (
-  dvcRoot: string,
-  experiments: Experiments
+  disposer: Disposer,
+  dvcRoot = dvcDemoPath
 ) => {
+  const { dvcExecutor, dvcRunner, experiments, experimentsModel, messageSpy } =
+    buildExperiments({ disposer })
+
   const mockGetOnlyOrPickProject = stub(
     WorkspaceExperiments.prototype,
     'getOnlyOrPickProject'
@@ -257,5 +260,13 @@ export const stubWorkspaceExperimentsGetters = (
     'getRepository'
   ).returns(experiments)
 
-  return [mockGetOnlyOrPickProject, mockGetRepository]
+  return {
+    dvcExecutor,
+    dvcRunner,
+    experiments,
+    experimentsModel,
+    messageSpy,
+    mockGetOnlyOrPickProject,
+    mockGetRepository
+  }
 }