From 18b2e18857d489d277e379e062eda7b90cb6cefc Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Tue, 5 Sep 2023 12:28:19 +1000 Subject: [PATCH] Patch Studio API update timinig issue --- extension/src/experiments/index.ts | 58 +++++++++++++++- .../src/experiments/model/collect.test.ts | 6 +- extension/src/experiments/model/collect.ts | 2 +- extension/src/experiments/model/index.ts | 13 +++- extension/src/experiments/studio.ts | 4 ++ .../src/test/suite/experiments/index.test.ts | 68 ++++++++++++++++++- extension/src/test/suite/experiments/util.ts | 6 +- 7 files changed, 144 insertions(+), 13 deletions(-) diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 579d5725bd..5cf30377c7 100644 --- a/extension/src/experiments/index.ts +++ b/extension/src/experiments/index.ts @@ -10,6 +10,7 @@ import omit from 'lodash.omit' import { ColumnLike, addStarredToColumns } from './columns/like' import { setContextForEditorTitleIcons } from './context' import { ExperimentsModel } from './model' +import { collectRemoteExpShas } from './model/collect' import { pickExperiment, pickExperiments, @@ -229,9 +230,10 @@ export class Experiments extends BaseRepository { } public async unsetPushing(ids: string[]) { - await this.update() + await Promise.all([this.update(), this.patchStudioApiTimingIssue(ids)]) + this.experiments.unsetPushing(ids) - return this.notifyChanged() + return this.webviewMessages.sendWebviewMessage() } public hasCheckpoints() { @@ -759,4 +761,56 @@ export class Experiments extends BaseRepository { const columnLikes = addStarredToColumns(columns) return pickColumnToFilter(columnLikes) } + + private async patchStudioApiTimingIssue(ids: string[]) { + if (!this.studio.isConnected()) { + return + } + + const [, { lsRemoteOutput }] = await Promise.all([ + this.waitForStudioUpdate(), + this.waitForRemoteUpdate() + ]) + + const remoteExpShas = collectRemoteExpShas(lsRemoteOutput) + + const shas = [] + for (const { id, sha } of this.experiments.getExperiments()) { + if (ids.includes(id) && sha && remoteExpShas.has(sha)) { + shas.push(sha) + } + } + + this.experiments.assumePushed(shas) + } + + private waitForStudioUpdate() { + return new Promise(resolve => { + const listener = this.dispose.track( + this.data.onDidUpdate(data => { + if (!isStudioExperimentsOutput(data)) { + return + } + this.dispose.untrack(listener) + listener.dispose() + resolve(undefined) + }) + ) + }) + } + + private waitForRemoteUpdate(): Promise<{ lsRemoteOutput: string }> { + return new Promise(resolve => { + const listener = this.dispose.track( + this.data.onDidUpdate(data => { + if (!isRemoteExperimentsOutput(data)) { + return + } + this.dispose.untrack(listener) + listener.dispose() + resolve(data) + }) + ) + }) + } } diff --git a/extension/src/experiments/model/collect.test.ts b/extension/src/experiments/model/collect.test.ts index 605dad0a26..7f74a472cc 100644 --- a/extension/src/experiments/model/collect.test.ts +++ b/extension/src/experiments/model/collect.test.ts @@ -1,4 +1,4 @@ -import { collectExperiments, collectRemoteExpRefs } from './collect' +import { collectExperiments, collectRemoteExpShas } from './collect' import { generateTestExpShowOutput } from '../../test/util/experiments' import { ExpShowOutput } from '../../cli/dvc/contract' @@ -102,7 +102,7 @@ describe('collectExperiments', () => { }) }) -describe('collectRemoteExpRefs', () => { +describe('collectRemoteExpShas', () => { it('should parse the git ls-remote output', () => { const output = `263e4408e42a0e215b0f70b36b2ab7b65a160d7e refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/vital-taal d4f2a35773ead55b7ce4b596f600e98360e49372 refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/whole-bout @@ -111,7 +111,7 @@ describe('collectRemoteExpRefs', () => { 390aef747f45fc49ec8928b24771f8950d057393 refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/known-flus 142a803b83ff784ba1106cc4ad0ba03310da6186 refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/tight-lira 21ce298cd1743405a0d73f5cb4cf52289ffa3276 refs/exps/bf/6ca8a35911bc6e62fb9bcaa506d4f4e185450c/crumb-orcs` - const remoteExpShas = collectRemoteExpRefs(output) + const remoteExpShas = collectRemoteExpShas(output) expect(remoteExpShas).toStrictEqual( new Set([ '263e4408e42a0e215b0f70b36b2ab7b65a160d7e', diff --git a/extension/src/experiments/model/collect.ts b/extension/src/experiments/model/collect.ts index 70848816ed..35990e909e 100644 --- a/extension/src/experiments/model/collect.ts +++ b/extension/src/experiments/model/collect.ts @@ -485,7 +485,7 @@ export const collectRunningInWorkspace = ( } } -export const collectRemoteExpRefs = (lsRemoteOutput: string): Set => { +export const collectRemoteExpShas = (lsRemoteOutput: string): Set => { const acc = new Set() for (const shaAndRef of trimAndSplit(lsRemoteOutput)) { const [sha] = shaAndRef.split(/\s+/) diff --git a/extension/src/experiments/model/index.ts b/extension/src/experiments/model/index.ts index 30a018c13e..1cbfb35d49 100644 --- a/extension/src/experiments/model/index.ts +++ b/extension/src/experiments/model/index.ts @@ -6,7 +6,7 @@ import { collectAddRemoveCommitsDetails, collectExperiments, collectOrderedCommitsAndExperiments, - collectRemoteExpRefs, + collectRemoteExpShas, collectRunningInQueue, collectRunningInWorkspace } from './collect' @@ -181,7 +181,7 @@ export class ExperimentsModel extends ModelWithPersistence { } public transformAndSetRemote(lsRemoteOutput: string) { - const remoteExpShas = collectRemoteExpRefs(lsRemoteOutput) + const remoteExpShas = collectRemoteExpShas(lsRemoteOutput) this.remoteExpShas = remoteExpShas } @@ -563,6 +563,15 @@ export class ExperimentsModel extends ModelWithPersistence { this.studioPushedExperiments = pushed } + public assumePushed(shas: string[]) { + for (const sha of shas) { + if (this.studioPushedExperiments.includes(sha)) { + continue + } + this.studioPushedExperiments.push(sha) + } + } + public hasDvcLiveOnlyRunning() { return !!this.dvcLiveOnlyExpName } diff --git a/extension/src/experiments/studio.ts b/extension/src/experiments/studio.ts index 8a6da91675..208adecbd7 100644 --- a/extension/src/experiments/studio.ts +++ b/extension/src/experiments/studio.ts @@ -36,6 +36,10 @@ export class Studio extends DeferredDisposable { return this.accessTokenSet } + public isConnected() { + return !!this.baseUrl + } + public getAccessToken() { return this.studioAccessToken } diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 1309677261..6c8c749023 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -13,7 +13,8 @@ import { CancellationToken, WorkspaceConfiguration, MessageItem, - ConfigurationTarget + ConfigurationTarget, + EventEmitter } from 'vscode' import { DEFAULT_EXPERIMENTS_OUTPUT, @@ -94,6 +95,7 @@ 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 { ExperimentsOutput } from '../../../data' const { openFileInEditor } = FileSystem @@ -593,8 +595,18 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should handle a message to push an experiment', async () => { - const { experiments, mockMessageReceived } = - await stubWorkspaceGettersWebview(disposable) + const { + experiments, + experimentsModel, + messageSpy, + mockMessageReceived, + webview + } = await stubWorkspaceGettersWebview(disposable) + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const studio = (experiments as any).studio + + const mockIsConnected = stub(studio, 'isConnected').returns(false) const mockExpId = 'exp-e7a67' @@ -634,6 +646,15 @@ suite('Experiments Test Suite', () => { RegisteredCommands.SETUP_SHOW_STUDIO_CONNECT ) + const experimentWithoutLink = experimentsModel + .getRowData()[1] + .subRows?.find(({ id }) => id === mockExpId) + + expect(experimentWithoutLink?.studioLinkType).not.to.equal( + StudioLinkType.PUSHED + ) + + mockIsConnected.restore() mockGetStudioAccessToken.resetBehavior() const tokenFound = new Promise(resolve => @@ -664,6 +685,25 @@ suite('Experiments Test Suite', () => { }) ) + const dataUpdated = disposable.track( + new EventEmitter() + ) + // eslint-disable-next-line @typescript-eslint/no-explicit-any + ;(experiments as any).data.onDidUpdate = dataUpdated.event + + let calls = 0 + + const remoteUpdated = new Promise(resolve => + disposable.track( + dataUpdated.event(() => { + calls = calls + 1 + if (calls === 2) { + resolve(undefined) + } + }) + ) + ) + mockMessageReceived.fire({ payload: [mockExpId], type: MessageFromWebviewType.PUSH_EXPERIMENT @@ -677,6 +717,28 @@ suite('Experiments Test Suite', () => { message: "Experiment major-lamb is up to date on Git remote 'origin'.\nView your experiments in [Studio](https://studio.iterative.ai/user/mattseddon/projects/vscode-dvc-demo-ynm6t3jxdx)" }) + + messageSpy.restore() + const mockShow = stub(webview, 'show') + + const messageSent = new Promise(resolve => + mockShow.callsFake(() => { + resolve(undefined) + return Promise.resolve(true) + }) + ) + dataUpdated.fire({ live: [], pushed: [], baseUrl: '' }) + dataUpdated.fire({ + lsRemoteOutput: `42b8736b08170529903cd203a1f40382a4b4a8cd refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/test-branch + 4fb124aebddb2adf1545030907687fa9a4c80e70 refs/exps/a9/53c3851f46955fa3e2b8f6e1c52999acc8c9ea77/${mockExpId}` + }) + await Promise.all([remoteUpdated, messageSent]) + + const experimentWithLink = experimentsModel + .getRowData()[1] + .subRows?.find(({ id }) => id === mockExpId) + + expect(experimentWithLink?.studioLinkType).to.equal(StudioLinkType.PUSHED) }).timeout(WEBVIEW_TEST_TIMEOUT) it("should handle a message to copy an experiment's Studio link", async () => { diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index db4b916d75..1b3ed1c617 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -384,7 +384,8 @@ export const stubWorkspaceGettersWebview = async ( experiments, experimentsModel, messageSpy, - mockMessageReceived + mockMessageReceived, + webview } = await buildExperimentsWebview({ disposer }) return { @@ -395,6 +396,7 @@ export const stubWorkspaceGettersWebview = async ( experimentsModel, messageSpy, ...stubWorkspaceExperiments(dvcRoot, experiments), - mockMessageReceived + mockMessageReceived, + webview } }