fix integration tests

extension/src/experiments/data/index.ts

@@ -14,7 +14,8 @@ import { ExpShowOutput } from '../../cli/dvc/contract'
 import {
   BaseData,
   ExperimentsOutput,
-  isRemoteExperimentsOutput
+  isRemoteExperimentsOutput,
+  isStudioExperimentsOutput
 } from '../../data'
 import {
   Args,
@@ -116,23 +117,24 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   }
 
   private async doStudio(shas: string[]) {
-    // return if studio not ready yet as we will call when it becomes ready
     await this.studio.isReady()
 
     const studioAccessToken = this.studio.getAccessToken()
 
     if (!studioAccessToken) {
       this.notifyChanged({ baseUrl: null, live: [], pushed: [] })
+      return
     }
 
     const gitRemoteUrl = this.studio.getGitRemoteUrl()
 
     if (shas.length === 0) {
       this.notifyChanged({ live: [], pushed: [] })
+      return
     }
 
     const params = querystring.stringify({
-      commits: [...shas, 'a9d8057e088d46842f15c3b6d1bb2e4befd5f677'],
+      commits: shas,
       git_remote_url: gitRemoteUrl
     })
 
@@ -274,8 +276,10 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   private waitForInitialLocalData() {
     const waitForInitialData = this.dispose.track(
       this.onDidUpdate(data => {
-        if (isRemoteExperimentsOutput(data)) {
-          // return early if Studio
+        if (
+          isRemoteExperimentsOutput(data) ||
+          isStudioExperimentsOutput(data)
+        ) {
           return
         }
         this.dispose.untrack(waitForInitialData)

extension/src/experiments/index.ts

@@ -612,11 +612,13 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public setStudioAccessToken(studioAccessToken: string | undefined) {
-    if (this.studio.getAccessToken() === studioAccessToken) {
+    const oldAccessToken = this.studio.getAccessToken()
+    const accessTokenInitialized = this.studio.isAccessTokenSet()
+    this.studio.setAccessToken(studioAccessToken)
+
+    if (!accessTokenInitialized || oldAccessToken === studioAccessToken) {
       return
     }
-
-    this.studio.setAccessToken(studioAccessToken)
     return this.data.managedUpdate()
   }
 

extension/src/experiments/studio.ts

@@ -10,17 +10,16 @@ export class Studio extends DeferredDisposable {
   private studioAccessToken: string | undefined
   private gitRemoteUrl?: string
 
+  private accessTokenSet = false
+
   constructor(dvcRoot: string, internalCommands: InternalCommands) {
     super()
     this.dvcRoot = dvcRoot
     this.internalCommands = internalCommands
 
     void this.internalCommands
       .executeCommand(AvailableCommands.GIT_GET_REMOTE_URL, this.dvcRoot)
-      .then(gitRemoteUrl => {
-        this.gitRemoteUrl = gitRemoteUrl
-        this.deferred.resolve()
-      })
+      .then(gitRemoteUrl => (this.gitRemoteUrl = gitRemoteUrl))
   }
 
   public getGitRemoteUrl() {
@@ -29,6 +28,12 @@ export class Studio extends DeferredDisposable {
 
   public setAccessToken(studioAccessToken: string | undefined) {
     this.studioAccessToken = studioAccessToken
+    this.accessTokenSet = true
+    this.deferred.resolve()
+  }
+
+  public isAccessTokenSet() {
+    return this.accessTokenSet
   }
 
   public getAccessToken() {

extension/src/test/suite/experiments/data/index.test.ts

@@ -28,7 +28,10 @@ import * as FileSystem from '../../../../fileSystem'
 import { ExperimentsModel } from '../../../../experiments/model'
 import { EXPERIMENT_WORKSPACE_ID } from '../../../../cli/dvc/contract'
 import expShowFixture from '../../../fixtures/expShow/base/output'
-import { isRemoteExperimentsOutput } from '../../../../data'
+import {
+  isRemoteExperimentsOutput,
+  isStudioExperimentsOutput
+} from '../../../../data'
 import { Studio } from '../../../../experiments/studio'
 
 const MOCK_WORKSPACE_GIT_FOLDER = join(dvcDemoPath, '.mock-git')
@@ -53,7 +56,10 @@ suite('Experiments Data Test Suite', () => {
     const getDataUpdatedEvent = (data: ExperimentsData) =>
       new Promise(resolve =>
         data.onDidUpdate(data => {
-          if (isRemoteExperimentsOutput(data)) {
+          if (
+            isRemoteExperimentsOutput(data) ||
+            isStudioExperimentsOutput(data)
+          ) {
             return
           }
 
@@ -127,7 +133,12 @@ suite('Experiments Data Test Suite', () => {
             }),
             setBranches: stub()
           } as unknown as ExperimentsModel,
-          { getAccessToken: () => Promise.resolve('') } as unknown as Studio,
+          {
+            getAccessToken: () => Promise.resolve(''),
+            getGitRemoteUrl: () =>
+              Promise.resolve('[email protected]:iterative/vscode-dvc-demo.git'),
+            isReady: () => Promise.resolve(undefined)
+          } as unknown as Studio,
           []
         )
       )
@@ -189,7 +200,12 @@ suite('Experiments Data Test Suite', () => {
             }),
             setBranches: stub()
           } as unknown as ExperimentsModel,
-          { getAccessToken: () => Promise.resolve('') } as unknown as Studio,
+          {
+            getAccessToken: () => Promise.resolve(''),
+            getGitRemoteUrl: () =>
+              Promise.resolve('[email protected]:iterative/vscode-dvc-demo.git'),
+            isReady: () => Promise.resolve(undefined)
+          } as unknown as Studio,
           []
         )
       )

extension/src/test/suite/experiments/index.test.ts

@@ -1,7 +1,6 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix */
 import { join, resolve } from 'path'
 import { after, afterEach, beforeEach, describe, it, suite } from 'mocha'
-import * as Fetch from 'node-fetch'
 import { expect } from 'chai'
 import { stub, spy, restore, SinonStub } from 'sinon'
 import {
@@ -95,7 +94,6 @@ import { MAX_SELECTED_EXPERIMENTS } from '../../../experiments/model/status'
 import { Pipeline } from '../../../pipeline'
 import { ColumnLike } from '../../../experiments/columns/like'
 import * as Clipboard from '../../../vscode/clipboard'
-import { STUDIO_URL } from '../../../setup/webview/contract'
 
 const { openFileInEditor } = FileSystem
 
@@ -682,23 +680,10 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it("should handle a message to copy an experiment's Studio link", async () => {
-      const { mockMessageReceived, experiments, gitReader } =
-        await buildExperimentsWebview({ disposer: disposable })
-
-      const viewUrl =
-        'https://studio.iterative.ai/user/demo-user/projects/demo-ynm6t3jxdx'
-
-      const mockGitRemoteUrl = '[email protected]:iterative/vscode-dvc-demo.git'
-      stub(gitReader, 'getRemoteUrl').resolves(mockGitRemoteUrl)
-      const mockStudioToken = 'isat_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB'
-      const mockFetch = stub(Fetch, 'default').resolves({
-        json: () =>
-          Promise.resolve({
-            url: viewUrl
-          })
-      } as unknown as Fetch.Response)
+      const { mockMessageReceived } = await buildExperimentsWebview({
+        disposer: disposable
+      })
 
-      await experiments.setStudioAccessToken(mockStudioToken)
       const mockWriteToClipboard = stub(Clipboard, 'writeToClipboard')
       const writeToClipboardCalled = new Promise(resolve =>
         mockWriteToClipboard.callsFake(() => {
@@ -712,36 +697,18 @@ suite('Experiments Test Suite', () => {
         payload: { id: 'exp-e7a67', type: StudioLinkType.PUSHED }
       })
 
-      expect(mockFetch).to.be.calledWith(`${STUDIO_URL}/webhook/dvc`, {
-        body: JSON.stringify({
-          client: 'vscode',
-          refs: {
-            pushed: [
-              'refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/test-branch'
-            ]
-          },
-          repo_url: mockGitRemoteUrl
-        }),
-        headers: {
-          Authorization: `token ${mockStudioToken}`,
-          'Content-Type': 'application/json'
-        },
-        method: 'POST'
-      })
-
       await writeToClipboardCalled
       const link =
-        viewUrl +
+        'https://studio.iterative.ai/user/olivaw/projects/vscode-dvc-demo-ynm6t3jxdx' +
         '?showOnlySelected=1' +
-        '&experimentReferences=4fb124aebddb2adf1545030907687fa9a4c80e70' +
-        '&activeExperimentReferences=4fb124aebddb2adf1545030907687fa9a4c80e70%3Aprimary'
+        '&experimentReferences=4fb124aebddb2adf1545030907687fa9a4c80e70'
 
       expect(mockWriteToClipboard).to.be.calledOnce
       expect(mockWriteToClipboard).to.be.calledWithExactly(
         link,
         `[Studio link](${link})`
       )
-    })
+    }).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, mockMessageReceived } =

extension/src/test/suite/experiments/util.ts

@@ -35,20 +35,26 @@ export const DEFAULT_EXPERIMENTS_OUTPUT = {
 
 export const buildExperiments = ({
   availableNbCommits = { main: 5 },
+  baseUrl = 'https://studio.iterative.ai/user/olivaw/projects/vscode-dvc-demo-ynm6t3jxdx',
   disposer,
   dvcRoot = dvcDemoPath,
   expShow = expShowFixture,
   gitLog = gitLogFixture,
+  live = [],
   lsRemoteOutput = remoteExpRefsFixture,
+  pushed = ['42b8736b08170529903cd203a1f40382a4b4a8cd'],
   rowOrder = rowOrderFixture,
   stageList = 'train'
 }: {
   availableNbCommits?: { [branch: string]: number }
   disposer: Disposer
+  baseUrl?: string
   dvcRoot?: string
   expShow?: ExpShowOutput
   gitLog?: string
+  live?: { baseline_sha: string; name: string }[]
   lsRemoteOutput?: string
+  pushed?: string[]
   rowOrder?: { branch: string; sha: string }[]
   stageList?: string | null
 }) => {
@@ -104,7 +110,8 @@ export const buildExperiments = ({
       gitLog,
       rowOrder
     }),
-    experiments.setState({ lsRemoteOutput })
+    experiments.setState({ lsRemoteOutput }),
+    experiments.setState({ baseUrl, live, pushed })
   ])
 
   return {
@@ -286,7 +293,12 @@ export const buildExperimentsData = (
         getNbOfCommitsToShow: () => DEFAULT_CURRENT_BRANCH_COMMITS_TO_SHOW,
         setBranches: mockSetBranches
       } as unknown as ExperimentsModel,
-      { getAccessToken: () => Promise.resolve('') } as unknown as Studio,
+      {
+        getAccessToken: () => Promise.resolve(''),
+        getGitRemoteUrl: () =>
+          Promise.resolve('[email protected]:iterative/vscode-dvc-demo.git'),
+        isReady: () => Promise.resolve(undefined)
+      } as unknown as Studio,
       []
     )
   )