diff --git a/extension/src/cli/dvc/runner.ts b/extension/src/cli/dvc/runner.ts index 9dbc9689ba..c1002e3024 100644 --- a/extension/src/cli/dvc/runner.ts +++ b/extension/src/cli/dvc/runner.ts @@ -46,16 +46,15 @@ export class DvcRunner extends Disposable implements ICli { private readonly pseudoTerminal: PseudoTerminal private currentProcess: Process | undefined private readonly config: Config - private readonly getStudioAccessToken: () => Thenable + private readonly getStudioLiveShareToken: () => string | undefined constructor( config: Config, - getStudioAccessToken: () => Thenable + getStudioLiveShareToken: () => string | undefined ) { super() - this.config = config - this.getStudioAccessToken = getStudioAccessToken + this.getStudioLiveShareToken = getStudioLiveShareToken this.processCompleted = this.dispose.track(new EventEmitter()) this.onDidCompleteProcess = this.processCompleted.event @@ -146,16 +145,8 @@ export class DvcRunner extends Disposable implements ICli { return this.currentProcess } - private createProcess({ - cwd, - args, - studioAccessToken - }: { - cwd: string - args: Args - studioAccessToken: string | undefined - }): Process { - const options = this.getOptions(cwd, args, studioAccessToken) + private createProcess({ cwd, args }: { cwd: string; args: Args }): Process { + const options = this.getOptions(cwd, args) const command = getCommandString(options) const stopWatch = new StopWatch() const process = this.dispose.track(createProcess(options)) @@ -183,11 +174,7 @@ export class DvcRunner extends Disposable implements ICli { return process } - private getOptions( - cwd: string, - args: Args, - studioAccessToken: string | undefined - ) { + private getOptions(cwd: string, args: Args) { const options = getOptions( this.config.getPythonBinPath(), this.config.getCliPath(), @@ -195,6 +182,8 @@ export class DvcRunner extends Disposable implements ICli { ...args ) + const studioAccessToken = this.getStudioLiveShareToken() + if (!studioAccessToken) { return options } @@ -205,18 +194,13 @@ export class DvcRunner extends Disposable implements ICli { } } - private async startProcess(cwd: string, args: Args) { + private startProcess(cwd: string, args: Args) { this.pseudoTerminal.setBlocked(true) this.processOutput.fire(`Running: dvc ${args.join(' ')}\r\n\n`) - const studioAccessToken = this.config.sendLiveToStudio() - ? await this.getStudioAccessToken() - : undefined - this.currentProcess = this.createProcess({ args, - cwd, - studioAccessToken + cwd }) } } diff --git a/extension/src/config.ts b/extension/src/config.ts index dca6dcf8e5..049c4fd04c 100644 --- a/extension/src/config.ts +++ b/extension/src/config.ts @@ -83,13 +83,6 @@ export class Config extends DeferredDisposable { return isPythonExtensionInstalled() } - public sendLiveToStudio() { - return getConfigValue( - ConfigKey.STUDIO_SHARE_EXPERIMENTS_LIVE, - false - ) - } - private async getConfigOrExtensionPythonBinPath() { return getConfigValue(ConfigKey.PYTHON_PATH) || (await getPythonBinPath()) } diff --git a/extension/src/connect/index.ts b/extension/src/connect/index.ts index 54e92504aa..f726e07d0d 100644 --- a/extension/src/connect/index.ts +++ b/extension/src/connect/index.ts @@ -14,11 +14,13 @@ import { ContextKey, setContextValue } from '../vscode/context' import { RegisteredCommands } from '../commands/external' import { Modal } from '../vscode/modal' import { GLOBAL_WEBVIEW_DVCROOT } from '../webview/factory' +import { ConfigKey, getConfigValue } from '../vscode/config' export class Connect extends BaseRepository { public readonly viewKey = ViewKey.CONNECT private readonly secrets: SecretStorage + private studioAccessToken: string | undefined = undefined constructor(context: ExtensionContext, webviewIcon: Resource) { super(GLOBAL_WEBVIEW_DVCROOT, webviewIcon) @@ -31,13 +33,20 @@ export class Connect extends BaseRepository { ) ) - void this.setContext().then(() => this.deferred.resolve()) + void this.getSecret(STUDIO_ACCESS_TOKEN_KEY).then( + async studioAccessToken => { + this.studioAccessToken = studioAccessToken + await this.setContext() + this.deferred.resolve() + } + ) this.dispose.track( - context.secrets.onDidChange(e => { + context.secrets.onDidChange(async e => { if (e.key !== STUDIO_ACCESS_TOKEN_KEY) { return } + this.studioAccessToken = await this.getSecret(STUDIO_ACCESS_TOKEN_KEY) return this.setContext() }) ) @@ -62,8 +71,17 @@ export class Connect extends BaseRepository { return this.storeSecret(STUDIO_ACCESS_TOKEN_KEY, token) } + public getStudioLiveShareToken() { + return getConfigValue( + ConfigKey.STUDIO_SHARE_EXPERIMENTS_LIVE, + false + ) + ? this.getStudioAccessToken() + : undefined + } + public getStudioAccessToken() { - return this.getSecret(STUDIO_ACCESS_TOKEN_KEY) + return this.studioAccessToken } private handleMessageFromWebview(message: MessageFromWebview) { @@ -89,8 +107,8 @@ export class Connect extends BaseRepository { return openUrl(`${STUDIO_URL}/user/_/profile?section=accessToken`) } - private async setContext() { - const storedToken = await this.getStudioAccessToken() + private setContext() { + const storedToken = this.getStudioAccessToken() if (isStudioAccessToken(storedToken)) { if (this.deferred.state === 'resolved') { void Modal.showInformation( diff --git a/extension/src/experiments/commands/index.ts b/extension/src/experiments/commands/index.ts index ce14527c64..fd7b1b68e6 100644 --- a/extension/src/experiments/commands/index.ts +++ b/extension/src/experiments/commands/index.ts @@ -109,8 +109,8 @@ export const getShareExperimentAsCommitCommand = export const getShareExperimentToStudioCommand = (internalCommands: InternalCommands, connect: Connect) => - async ({ dvcRoot, id }: { dvcRoot: string; id: string }) => { - const studioAccessToken = await connect.getStudioAccessToken() + ({ dvcRoot, id }: { dvcRoot: string; id: string }) => { + const studioAccessToken = connect.getStudioAccessToken() if (!studioAccessToken) { return commands.executeCommand(RegisteredCommands.CONNECT_SHOW) } diff --git a/extension/src/extension.ts b/extension/src/extension.ts index c304c1c545..31c64bc0e7 100644 --- a/extension/src/extension.ts +++ b/extension/src/extension.ts @@ -97,7 +97,7 @@ export class Extension extends Disposable { this.dvcExecutor = this.dispose.track(new DvcExecutor(config)) this.dvcReader = this.dispose.track(new DvcReader(config)) this.dvcRunner = this.dispose.track( - new DvcRunner(config, () => this.connect.getStudioAccessToken()) + new DvcRunner(config, () => this.connect.getStudioLiveShareToken()) ) this.dvcViewer = this.dispose.track(new DvcViewer(config)) diff --git a/extension/src/test/suite/cli/dvc/runner.test.ts b/extension/src/test/suite/cli/dvc/runner.test.ts index 656ac39929..b103bd272f 100644 --- a/extension/src/test/suite/cli/dvc/runner.test.ts +++ b/extension/src/test/suite/cli/dvc/runner.test.ts @@ -26,13 +26,12 @@ suite('DVC Runner Test Suite', () => { it('should only be able to run a single command at a time', async () => { const mockConfig = { getCliPath: () => 'sleep', - getPythonBinPath: () => undefined, - sendLiveToStudio: () => false + getPythonBinPath: () => undefined } as Config - const mockGetStudioAccessToken = () => Promise.resolve(undefined) + const mockGetStudioLiveShareToken = () => undefined const dvcRunner = disposable.track( - new DvcRunner(mockConfig, mockGetStudioAccessToken) + new DvcRunner(mockConfig, mockGetStudioLiveShareToken) ) const windowErrorMessageSpy = spy(window, 'showErrorMessage') @@ -47,13 +46,12 @@ suite('DVC Runner Test Suite', () => { it('should be able to stop a started command', async () => { const mockConfig = { getCliPath: () => 'sleep', - getPythonBinPath: () => undefined, - sendLiveToStudio: () => false + getPythonBinPath: () => undefined } as Config - const mockGetStudioAccessToken = () => Promise.resolve(undefined) + const mockGetStudioLiveShareToken = () => undefined const dvcRunner = disposable.track( - new DvcRunner(mockConfig, mockGetStudioAccessToken) + new DvcRunner(mockConfig, mockGetStudioLiveShareToken) ) const cwd = __dirname @@ -116,13 +114,12 @@ suite('DVC Runner Test Suite', () => { const mockConfig = { getCliPath: () => 'echo', - getPythonBinPath: () => undefined, - sendLiveToStudio: () => false + getPythonBinPath: () => undefined } as Config - const mockGetStudioAccessToken = () => Promise.resolve(undefined) + const mockGetStudioLiveShareToken = () => undefined const dvcRunner = disposable.track( - new DvcRunner(mockConfig, mockGetStudioAccessToken) + new DvcRunner(mockConfig, mockGetStudioLiveShareToken) ) const started = onDidStartOrCompleteProcess(dvcRunner.onDidStartProcess) @@ -155,17 +152,16 @@ suite('DVC Runner Test Suite', () => { const mockConfig = { getCliPath: () => 'echo', - getPythonBinPath: () => undefined, - sendLiveToStudio: () => true + getPythonBinPath: () => undefined } as Config - const mockGetStudioAccessToken = stub() + const mockGetStudioLiveShareToken = stub() .onFirstCall() - .resolves(mockStudioAccessToken) + .returns(mockStudioAccessToken) .onSecondCall() - .resolves(undefined) + .returns(undefined) const dvcRunner = disposable.track( - new DvcRunner(mockConfig, mockGetStudioAccessToken) + new DvcRunner(mockConfig, mockGetStudioLiveShareToken) ) await dvcRunner.runExperiment(dvcDemoPath) diff --git a/extension/src/test/suite/connect/index.test.ts b/extension/src/test/suite/connect/index.test.ts index 69421256a5..a0d8523a07 100644 --- a/extension/src/test/suite/connect/index.test.ts +++ b/extension/src/test/suite/connect/index.test.ts @@ -40,7 +40,10 @@ suite('Connect Test Suite', () => { const connect = disposable.track( new Connect( { - secrets: mockSecretStorage || { get: stub(), onDidChange: stub() } + secrets: mockSecretStorage || { + get: stub().resolves(undefined), + onDidChange: stub() + } } as unknown as ExtensionContext, resourceLocator.dvcIcon ) diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 0818bdc1db..19478e4925 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -608,7 +608,7 @@ suite('Experiments Test Suite', () => { const tokenAccessed = new Promise(resolve => mockGetStudioAccessToken.callsFake(() => { resolve(undefined) - return Promise.resolve(undefined) + return undefined }) ) diff --git a/extension/src/test/suite/util.ts b/extension/src/test/suite/util.ts index 250630d004..d6bab180d6 100644 --- a/extension/src/test/suite/util.ts +++ b/extension/src/test/suite/util.ts @@ -140,9 +140,7 @@ export const bypassProcessManagerDebounce = ( export const buildInternalCommands = (disposer: Disposer) => { const config = disposer.track(new Config()) const dvcReader = disposer.track(new DvcReader(config)) - const dvcRunner = disposer.track( - new DvcRunner(config, () => Promise.resolve(undefined)) - ) + const dvcRunner = disposer.track(new DvcRunner(config, () => undefined)) const dvcExecutor = disposer.track(new DvcExecutor(config)) const dvcViewer = disposer.track(new DvcViewer(config)) const gitReader = disposer.track(new GitReader())