From f43826256703a40f76e1a93e677e72e5963689bc Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Tue, 17 Oct 2023 19:38:08 +1100 Subject: [PATCH] Add support for a tensorboard experiment (#22215) --- package.json | 6 +- package.nls.json | 1 + src/client/common/experiments/groups.ts | 5 + .../nbextensionCodeLensProvider.ts | 22 +- src/client/tensorBoard/serviceRegistry.ts | 2 + .../tensorBoard/tensorBoardFileWatcher.ts | 22 +- .../tensorBoardImportCodeLensProvider.ts | 22 +- .../tensorBoard/tensorBoardSessionProvider.ts | 28 ++- .../tensorBoard/tensorBoardUsageTracker.ts | 16 +- .../tensorBoard/tensorboarExperiment.ts | 65 +++++- .../tensorBoard/tensorboardIntegration.ts | 3 - src/client/tensorBoard/terminalWatcher.ts | 12 +- .../nbextensionCodeLensProvider.unit.test.ts | 85 ++++---- ...orBoardImportCodeLensProvider.unit.test.ts | 106 +++++----- .../tensorBoardPrompt.unit.test.ts | 2 +- .../tensorBoardUsageTracker.unit.test.ts | 191 ++++++++++-------- 16 files changed, 380 insertions(+), 208 deletions(-) diff --git a/package.json b/package.json index 20401cc43762..5b13f9eae0a3 100644 --- a/package.json +++ b/package.json @@ -537,7 +537,8 @@ "pythonPromptNewToolsExt", "pythonTerminalEnvVarActivation", "pythonTestAdapter", - "pythonREPLSmartSend" + "pythonREPLSmartSend", + "pythonRecommendTensorboardExt" ], "enumDescriptions": [ "%python.experiments.All.description%", @@ -545,7 +546,8 @@ "%python.experiments.pythonPromptNewToolsExt.description%", "%python.experiments.pythonTerminalEnvVarActivation.description%", "%python.experiments.pythonTestAdapter.description%", - "%python.experiments.pythonREPLSmartSend.description%" + "%python.experiments.pythonREPLSmartSend.description%", + "%python.experiments.pythonRecommendTensorboardExt.description%" ] }, "scope": "machine", diff --git a/package.nls.json b/package.nls.json index f843399e09c5..87692fb7c1a8 100644 --- a/package.nls.json +++ b/package.nls.json @@ -42,6 +42,7 @@ "python.experiments.pythonTerminalEnvVarActivation.description": "Enables use of environment variables to activate terminals instead of sending activation commands.", "python.experiments.pythonTestAdapter.description": "Denotes the Python Test Adapter experiment.", "python.experiments.pythonREPLSmartSend.description": "Denotes the Python REPL Smart Send experiment.", + "python.experiments.pythonRecommendTensorboardExt.description": "Denotes the Tensorboard Extension recommendation experiment.", "python.globalModuleInstallation.description": "Whether to install Python modules globally when not using an environment.", "python.languageServer.description": "Defines type of the language server.", "python.languageServer.defaultDescription": "Automatically select a language server: Pylance if installed and available, otherwise fallback to Jedi.", diff --git a/src/client/common/experiments/groups.ts b/src/client/common/experiments/groups.ts index b7a598e0a08a..29035bbc57fe 100644 --- a/src/client/common/experiments/groups.ts +++ b/src/client/common/experiments/groups.ts @@ -22,3 +22,8 @@ export enum EnableTestAdapterRewrite { export enum EnableREPLSmartSend { experiment = 'pythonREPLSmartSend', } + +// Experiment to recommend installing the tensorboard extension. +export enum RecommendTensobardExtension { + experiment = 'pythonRecommendTensorboardExt', +} diff --git a/src/client/tensorBoard/nbextensionCodeLensProvider.ts b/src/client/tensorBoard/nbextensionCodeLensProvider.ts index 7b9a116ee144..afaaf116851a 100644 --- a/src/client/tensorBoard/nbextensionCodeLensProvider.ts +++ b/src/client/tensorBoard/nbextensionCodeLensProvider.ts @@ -3,21 +3,23 @@ import { inject, injectable } from 'inversify'; import { once } from 'lodash'; -import { CancellationToken, CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode'; +import { CancellationToken, CodeLens, Command, Disposable, languages, Position, Range, TextDocument } from 'vscode'; import { IExtensionSingleActivationService } from '../activation/types'; import { Commands, NotebookCellScheme, PYTHON_LANGUAGE } from '../common/constants'; -import { IDisposableRegistry } from '../common/types'; +import { IDisposable, IDisposableRegistry } from '../common/types'; import { TensorBoard } from '../common/utils/localize'; import { sendTelemetryEvent } from '../telemetry'; import { EventName } from '../telemetry/constants'; import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants'; import { containsNotebookExtension } from './helpers'; -import { useNewTensorboardExtension } from './tensorboarExperiment'; +import { TensorboardExperiment } from './tensorboarExperiment'; @injectable() export class TensorBoardNbextensionCodeLensProvider implements IExtensionSingleActivationService { public readonly supportedWorkspaceTypes = { untrustedWorkspace: false, virtualWorkspace: false }; + private readonly disposables: IDisposable[] = []; + private sendTelemetryOnce = once( sendTelemetryEvent.bind(this, EventName.TENSORBOARD_ENTRYPOINT_SHOWN, undefined, { trigger: TensorBoardEntrypointTrigger.nbextension, @@ -25,12 +27,22 @@ export class TensorBoardNbextensionCodeLensProvider implements IExtensionSingleA }), ); - constructor(@inject(IDisposableRegistry) private disposables: IDisposableRegistry) {} + constructor( + @inject(IDisposableRegistry) disposables: IDisposableRegistry, + @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment, + ) { + disposables.push(this); + } + + public dispose(): void { + Disposable.from(...this.disposables).dispose(); + } public async activate(): Promise { - if (useNewTensorboardExtension()) { + if (TensorboardExperiment.isTensorboardExtensionInstalled) { return; } + this.experiment.disposeOnInstallingTensorboard(this); this.activateInternal().ignoreErrors(); } diff --git a/src/client/tensorBoard/serviceRegistry.ts b/src/client/tensorBoard/serviceRegistry.ts index dd193f528eea..5fedb7b6abf5 100644 --- a/src/client/tensorBoard/serviceRegistry.ts +++ b/src/client/tensorBoard/serviceRegistry.ts @@ -11,6 +11,7 @@ import { TensorBoardSessionProvider } from './tensorBoardSessionProvider'; import { TensorBoardNbextensionCodeLensProvider } from './nbextensionCodeLensProvider'; import { TerminalWatcher } from './terminalWatcher'; import { TensorboardDependencyChecker } from './tensorboardDependencyChecker'; +import { TensorboardExperiment } from './tensorboarExperiment'; export function registerTypes(serviceManager: IServiceManager): void { serviceManager.addSingleton(TensorBoardSessionProvider, TensorBoardSessionProvider); @@ -34,4 +35,5 @@ export function registerTypes(serviceManager: IServiceManager): void { serviceManager.addBinding(TensorBoardNbextensionCodeLensProvider, IExtensionSingleActivationService); serviceManager.addSingleton(IExtensionSingleActivationService, TerminalWatcher); serviceManager.addSingleton(TensorboardDependencyChecker, TensorboardDependencyChecker); + serviceManager.addSingleton(TensorboardExperiment, TensorboardExperiment); } diff --git a/src/client/tensorBoard/tensorBoardFileWatcher.ts b/src/client/tensorBoard/tensorBoardFileWatcher.ts index dccdb95290ec..f2f9344d7365 100644 --- a/src/client/tensorBoard/tensorBoardFileWatcher.ts +++ b/src/client/tensorBoard/tensorBoardFileWatcher.ts @@ -2,13 +2,13 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; -import { FileSystemWatcher, RelativePattern, WorkspaceFolder, WorkspaceFoldersChangeEvent } from 'vscode'; +import { Disposable, FileSystemWatcher, RelativePattern, WorkspaceFolder, WorkspaceFoldersChangeEvent } from 'vscode'; import { IExtensionSingleActivationService } from '../activation/types'; import { IWorkspaceService } from '../common/application/types'; -import { IDisposableRegistry } from '../common/types'; +import { IDisposable, IDisposableRegistry } from '../common/types'; import { TensorBoardEntrypointTrigger } from './constants'; import { TensorBoardPrompt } from './tensorBoardPrompt'; -import { useNewTensorboardExtension } from './tensorboarExperiment'; +import { TensorboardExperiment } from './tensorboarExperiment'; @injectable() export class TensorBoardFileWatcher implements IExtensionSingleActivationService { @@ -18,16 +18,26 @@ export class TensorBoardFileWatcher implements IExtensionSingleActivationService private globPatterns = ['*tfevents*', '*/*tfevents*', '*/*/*tfevents*']; + private readonly disposables: IDisposable[] = []; + constructor( @inject(IWorkspaceService) private workspaceService: IWorkspaceService, @inject(TensorBoardPrompt) private tensorBoardPrompt: TensorBoardPrompt, - @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, - ) {} + @inject(IDisposableRegistry) disposables: IDisposableRegistry, + @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment, + ) { + disposables.push(this); + } + + public dispose(): void { + Disposable.from(...this.disposables).dispose(); + } public async activate(): Promise { - if (useNewTensorboardExtension()) { + if (TensorboardExperiment.isTensorboardExtensionInstalled) { return; } + this.experiment.disposeOnInstallingTensorboard(this); this.activateInternal().ignoreErrors(); } diff --git a/src/client/tensorBoard/tensorBoardImportCodeLensProvider.ts b/src/client/tensorBoard/tensorBoardImportCodeLensProvider.ts index d6dc8d7e82e5..585b9151922a 100644 --- a/src/client/tensorBoard/tensorBoardImportCodeLensProvider.ts +++ b/src/client/tensorBoard/tensorBoardImportCodeLensProvider.ts @@ -3,16 +3,16 @@ import { inject, injectable } from 'inversify'; import { once } from 'lodash'; -import { CancellationToken, CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode'; +import { CancellationToken, CodeLens, Command, Disposable, languages, Position, Range, TextDocument } from 'vscode'; import { IExtensionSingleActivationService } from '../activation/types'; import { Commands, PYTHON } from '../common/constants'; -import { IDisposableRegistry } from '../common/types'; +import { IDisposable, IDisposableRegistry } from '../common/types'; import { TensorBoard } from '../common/utils/localize'; import { sendTelemetryEvent } from '../telemetry'; import { EventName } from '../telemetry/constants'; import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants'; import { containsTensorBoardImport } from './helpers'; -import { useNewTensorboardExtension } from './tensorboarExperiment'; +import { TensorboardExperiment } from './tensorboarExperiment'; @injectable() export class TensorBoardImportCodeLensProvider implements IExtensionSingleActivationService { @@ -25,12 +25,24 @@ export class TensorBoardImportCodeLensProvider implements IExtensionSingleActiva }), ); - constructor(@inject(IDisposableRegistry) private disposables: IDisposableRegistry) {} + private readonly disposables: IDisposable[] = []; + + constructor( + @inject(IDisposableRegistry) disposables: IDisposableRegistry, + @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment, + ) { + disposables.push(this); + } + + public dispose(): void { + Disposable.from(...this.disposables).dispose(); + } public async activate(): Promise { - if (useNewTensorboardExtension()) { + if (TensorboardExperiment.isTensorboardExtensionInstalled) { return; } + this.experiment.disposeOnInstallingTensorboard(this); this.activateInternal().ignoreErrors(); } diff --git a/src/client/tensorBoard/tensorBoardSessionProvider.ts b/src/client/tensorBoard/tensorBoardSessionProvider.ts index c81059654075..ec52b9ef94dc 100644 --- a/src/client/tensorBoard/tensorBoardSessionProvider.ts +++ b/src/client/tensorBoard/tensorBoardSessionProvider.ts @@ -2,7 +2,7 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; -import { l10n, ViewColumn } from 'vscode'; +import { Disposable, l10n, ViewColumn } from 'vscode'; import { IExtensionSingleActivationService } from '../activation/types'; import { IApplicationShell, ICommandManager, IWorkspaceService } from '../common/application/types'; import { Commands } from '../common/constants'; @@ -14,6 +14,7 @@ import { IPersistentState, IPersistentStateFactory, IConfigurationService, + IDisposable, } from '../common/types'; import { IMultiStepInputFactory } from '../common/utils/multiStepInput'; import { IInterpreterService } from '../interpreter/contracts'; @@ -22,7 +23,7 @@ import { sendTelemetryEvent } from '../telemetry'; import { EventName } from '../telemetry/constants'; import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants'; import { TensorBoardSession } from './tensorBoardSession'; -import { useNewTensorboardExtension } from './tensorboarExperiment'; +import { TensorboardExperiment } from './tensorboarExperiment'; export const PREFERRED_VIEWGROUP = 'PythonTensorBoardWebviewPreferredViewGroup'; @@ -36,18 +37,22 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer private hasActiveTensorBoardSessionContext: ContextKey; + private readonly disposables: IDisposable[] = []; + constructor( @inject(IInstaller) private readonly installer: IInstaller, @inject(IInterpreterService) private readonly interpreterService: IInterpreterService, @inject(IApplicationShell) private readonly applicationShell: IApplicationShell, @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService, @inject(ICommandManager) private readonly commandManager: ICommandManager, - @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, + @inject(IDisposableRegistry) disposables: IDisposableRegistry, @inject(IPythonExecutionFactory) private readonly pythonExecFactory: IPythonExecutionFactory, @inject(IPersistentStateFactory) private stateFactory: IPersistentStateFactory, @inject(IMultiStepInputFactory) private readonly multiStepFactory: IMultiStepInputFactory, @inject(IConfigurationService) private readonly configurationService: IConfigurationService, + @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment, ) { + disposables.push(this); this.preferredViewGroupMemento = this.stateFactory.createGlobalPersistentState( PREFERRED_VIEWGROUP, ViewColumn.Active, @@ -58,10 +63,15 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer ); } + public dispose(): void { + Disposable.from(...this.disposables).dispose(); + } + public async activate(): Promise { - if (useNewTensorboardExtension()) { + if (TensorboardExperiment.isTensorboardExtensionInstalled) { return; } + this.experiment.disposeOnInstallingTensorboard(this); this.disposables.push( this.commandManager.registerCommand( @@ -69,16 +79,20 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer ( entrypoint: TensorBoardEntrypoint = TensorBoardEntrypoint.palette, trigger: TensorBoardEntrypointTrigger = TensorBoardEntrypointTrigger.palette, - ) => { + ): void => { sendTelemetryEvent(EventName.TENSORBOARD_SESSION_LAUNCH, undefined, { trigger, entrypoint, }); - return this.createNewSession(); + if (this.experiment.recommendAndUseNewExtension() === 'continueWithPythonExtension') { + void this.createNewSession(); + } }, ), this.commandManager.registerCommand(Commands.RefreshTensorBoard, () => - this.knownSessions.map((w) => w.refresh()), + this.experiment.recommendAndUseNewExtension() === 'continueWithPythonExtension' + ? this.knownSessions.map((w) => w.refresh()) + : undefined, ), ); } diff --git a/src/client/tensorBoard/tensorBoardUsageTracker.ts b/src/client/tensorBoard/tensorBoardUsageTracker.ts index 7c8ea7b00961..b88e416a113f 100644 --- a/src/client/tensorBoard/tensorBoardUsageTracker.ts +++ b/src/client/tensorBoard/tensorBoardUsageTracker.ts @@ -3,7 +3,7 @@ import { inject, injectable } from 'inversify'; import * as path from 'path'; -import { TextEditor } from 'vscode'; +import { Disposable, TextEditor } from 'vscode'; import { IExtensionSingleActivationService } from '../activation/types'; import { IDocumentManager } from '../common/application/types'; import { isTestExecution } from '../common/constants'; @@ -12,7 +12,7 @@ import { getDocumentLines } from '../telemetry/importTracker'; import { TensorBoardEntrypointTrigger } from './constants'; import { containsTensorBoardImport } from './helpers'; import { TensorBoardPrompt } from './tensorBoardPrompt'; -import { useNewTensorboardExtension } from './tensorboarExperiment'; +import { TensorboardExperiment } from './tensorboarExperiment'; const testExecution = isTestExecution(); @@ -26,12 +26,20 @@ export class TensorBoardUsageTracker implements IExtensionSingleActivationServic @inject(IDocumentManager) private documentManager: IDocumentManager, @inject(IDisposableRegistry) private disposables: IDisposableRegistry, @inject(TensorBoardPrompt) private prompt: TensorBoardPrompt, - ) {} + @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment, + ) { + disposables.push(this); + } + + public dispose(): void { + Disposable.from(...this.disposables).dispose(); + } public async activate(): Promise { - if (useNewTensorboardExtension()) { + if (TensorboardExperiment.isTensorboardExtensionInstalled) { return; } + this.experiment.disposeOnInstallingTensorboard(this); if (testExecution) { await this.activateInternal(); } else { diff --git a/src/client/tensorBoard/tensorboarExperiment.ts b/src/client/tensorBoard/tensorboarExperiment.ts index 25eac8db71da..3cf4cb3c779a 100644 --- a/src/client/tensorBoard/tensorboarExperiment.ts +++ b/src/client/tensorBoard/tensorboarExperiment.ts @@ -1,8 +1,67 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { extensions } from 'vscode'; +import { Disposable, EventEmitter, commands, extensions, l10n, window } from 'vscode'; +import { inject, injectable } from 'inversify'; +import { IDisposable, IDisposableRegistry, IExperimentService } from '../common/types'; +import { RecommendTensobardExtension } from '../common/experiments/groups'; +import { TENSORBOARD_EXTENSION_ID } from '../common/constants'; -export function useNewTensorboardExtension(): boolean { - return !!extensions.getExtension('ms-toolsai.tensorboard'); +@injectable() +export class TensorboardExperiment { + private readonly _onDidChange = new EventEmitter(); + + public readonly onDidChange = this._onDidChange.event; + + private readonly toDisposeWhenTensobardIsInstalled: IDisposable[] = []; + + public static get isTensorboardExtensionInstalled(): boolean { + return !!extensions.getExtension(TENSORBOARD_EXTENSION_ID); + } + + private readonly isExperimentEnabled: boolean; + + constructor( + @inject(IDisposableRegistry) disposables: IDisposableRegistry, + @inject(IExperimentService) experiments: IExperimentService, + ) { + this.isExperimentEnabled = experiments.inExperimentSync(RecommendTensobardExtension.experiment); + disposables.push(this._onDidChange); + extensions.onDidChange( + () => + TensorboardExperiment.isTensorboardExtensionInstalled + ? Disposable.from(...this.toDisposeWhenTensobardIsInstalled).dispose() + : undefined, + this, + disposables, + ); + } + + public recommendAndUseNewExtension(): 'continueWithPythonExtension' | 'usingTensorboardExtension' { + if (!this.isExperimentEnabled) { + return 'continueWithPythonExtension'; + } + if (TensorboardExperiment.isTensorboardExtensionInstalled) { + return 'usingTensorboardExtension'; + } + const install = l10n.t('Install Tensorboard Extension'); + window + .showInformationMessage( + l10n.t( + 'Install the TensorBoard extension to use the this functionality. Once installed, select the command `Launch Tensorboard`.', + ), + { modal: true }, + install, + ) + .then((result): void => { + if (result === install) { + void commands.executeCommand('workbench.extensions.installExtension', TENSORBOARD_EXTENSION_ID); + } + }); + return 'usingTensorboardExtension'; + } + + public disposeOnInstallingTensorboard(disposabe: IDisposable): void { + this.toDisposeWhenTensobardIsInstalled.push(disposabe); + } } diff --git a/src/client/tensorBoard/tensorboardIntegration.ts b/src/client/tensorBoard/tensorboardIntegration.ts index 74f69afab84f..22d590d6ee65 100644 --- a/src/client/tensorBoard/tensorboardIntegration.ts +++ b/src/client/tensorBoard/tensorboardIntegration.ts @@ -69,10 +69,7 @@ export class TensorboardExtensionIntegration { public hideCommands(): void { if (this.extensions.getExtension(TENSORBOARD_EXTENSION_ID)) { - console.error('TensorBoard extension is installed'); void commands.executeCommand('setContext', 'python.tensorboardExtInstalled', true); - } else { - console.error('TensorBoard extension not installed'); } } diff --git a/src/client/tensorBoard/terminalWatcher.ts b/src/client/tensorBoard/terminalWatcher.ts index 30ccf7e1726a..5f48def54e43 100644 --- a/src/client/tensorBoard/terminalWatcher.ts +++ b/src/client/tensorBoard/terminalWatcher.ts @@ -4,7 +4,7 @@ import { IExtensionSingleActivationService } from '../activation/types'; import { IDisposable, IDisposableRegistry } from '../common/types'; import { sendTelemetryEvent } from '../telemetry'; import { EventName } from '../telemetry/constants'; -import { useNewTensorboardExtension } from './tensorboarExperiment'; +import { TensorboardExperiment } from './tensorboarExperiment'; // Every 5 min look, through active terminals to see if any are running `tensorboard` @injectable() @@ -13,12 +13,18 @@ export class TerminalWatcher implements IExtensionSingleActivationService, IDisp private handle: NodeJS.Timeout | undefined; - constructor(@inject(IDisposableRegistry) private disposables: IDisposableRegistry) {} + constructor( + @inject(IDisposableRegistry) private disposables: IDisposableRegistry, + @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment, + ) { + disposables.push(this); + } public async activate(): Promise { - if (useNewTensorboardExtension()) { + if (TensorboardExperiment.isTensorboardExtensionInstalled) { return; } + this.experiment.disposeOnInstallingTensorboard(this); const handle = setInterval(() => { // When user runs a command in VSCode terminal, the terminal's name // becomes the program that is currently running. Since tensorboard diff --git a/src/test/tensorBoard/nbextensionCodeLensProvider.unit.test.ts b/src/test/tensorBoard/nbextensionCodeLensProvider.unit.test.ts index 9a46d92c1422..aef90d14eacf 100644 --- a/src/test/tensorBoard/nbextensionCodeLensProvider.unit.test.ts +++ b/src/test/tensorBoard/nbextensionCodeLensProvider.unit.test.ts @@ -1,49 +1,60 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +import * as sinon from 'sinon'; import { assert } from 'chai'; import { CancellationTokenSource } from 'vscode'; +import { instance, mock } from 'ts-mockito'; import { TensorBoardNbextensionCodeLensProvider } from '../../client/tensorBoard/nbextensionCodeLensProvider'; import { MockDocument } from '../mocks/mockDocument'; +import { TensorboardExperiment } from '../../client/tensorBoard/tensorboarExperiment'; -suite('TensorBoard nbextension code lens provider', () => { - let codeLensProvider: TensorBoardNbextensionCodeLensProvider; - let cancelTokenSource: CancellationTokenSource; +[true, false].forEach((tbExtensionInstalled) => { + suite(`Tensorboard Extension is ${tbExtensionInstalled ? 'installed' : 'not installed'}`, () => { + suite.only('TensorBoard nbextension code lens provider', () => { + let experiment: TensorboardExperiment; + let codeLensProvider: TensorBoardNbextensionCodeLensProvider; + let cancelTokenSource: CancellationTokenSource; - setup(() => { - codeLensProvider = new TensorBoardNbextensionCodeLensProvider([]); - cancelTokenSource = new CancellationTokenSource(); - }); - teardown(() => { - cancelTokenSource.dispose(); - }); + setup(() => { + sinon.stub(TensorboardExperiment, 'isTensorboardExtensionInstalled').returns(tbExtensionInstalled); + experiment = mock(); + codeLensProvider = new TensorBoardNbextensionCodeLensProvider([], instance(experiment)); + cancelTokenSource = new CancellationTokenSource(); + }); + teardown(() => { + sinon.restore(); + cancelTokenSource.dispose(); + }); - test('Provide code lens for Python notebook loading tensorboard nbextension', async () => { - const document = new MockDocument('a=1\n%load_ext tensorboard', 'foo.ipynb', async () => true); - const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); - assert.ok(codeLens.length > 0, 'Failed to provide code lens for file loading tensorboard nbextension'); - }); - test('Provide code lens for Python notebook launching tensorboard nbextension', async () => { - const document = new MockDocument('a=1\n%tensorboard --logdir logs/fit', 'foo.ipynb', async () => true); - const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); - assert.ok(codeLens.length > 0, 'Failed to provide code lens for file loading tensorboard nbextension'); - }); - test('Fails when cancellation is signaled', () => { - const document = new MockDocument('a=1\n%tensorboard --logdir logs/fit', 'foo.ipynb', async () => true); - cancelTokenSource.cancel(); - const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); - assert.ok(codeLens.length === 0, 'Provided codelens even after cancellation was requested'); + test('Provide code lens for Python notebook loading tensorboard nbextension', async () => { + const document = new MockDocument('a=1\n%load_ext tensorboard', 'foo.ipynb', async () => true); + const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); + assert.ok(codeLens.length > 0, 'Failed to provide code lens for file loading tensorboard nbextension'); + }); + test('Provide code lens for Python notebook launching tensorboard nbextension', async () => { + const document = new MockDocument('a=1\n%tensorboard --logdir logs/fit', 'foo.ipynb', async () => true); + const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); + assert.ok(codeLens.length > 0, 'Failed to provide code lens for file loading tensorboard nbextension'); + }); + test('Fails when cancellation is signaled', () => { + const document = new MockDocument('a=1\n%tensorboard --logdir logs/fit', 'foo.ipynb', async () => true); + cancelTokenSource.cancel(); + const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); + assert.ok(codeLens.length === 0, 'Provided codelens even after cancellation was requested'); + }); + // Can't verify these cases without running in vscode as we depend on vscode to not call us + // based on the DocumentSelector we provided. See nbExtensionCodeLensProvider.test.ts for that. + // test('Does not provide code lens for Python file loading tensorboard nbextension', async () => { + // const document = new MockDocument('a=1\n%load_ext tensorboard', 'foo.py', async () => true); + // const codeLens = codeLensProvider.provideCodeLenses(document); + // assert.ok(codeLens.length === 0, 'Provided code lens for Python file loading tensorboard nbextension'); + // }); + // test('Does not provide code lens for Python file launching tensorboard nbextension', async () => { + // const document = new MockDocument('a=1\n%tensorboard --logdir logs/fit', 'foo.py', async () => true); + // const codeLens = codeLensProvider.provideCodeLenses(document); + // assert.ok(codeLens.length === 0, 'Provided code lens for Python file loading tensorboard nbextension'); + // }); + }); }); - // Can't verify these cases without running in vscode as we depend on vscode to not call us - // based on the DocumentSelector we provided. See nbExtensionCodeLensProvider.test.ts for that. - // test('Does not provide code lens for Python file loading tensorboard nbextension', async () => { - // const document = new MockDocument('a=1\n%load_ext tensorboard', 'foo.py', async () => true); - // const codeLens = codeLensProvider.provideCodeLenses(document); - // assert.ok(codeLens.length === 0, 'Provided code lens for Python file loading tensorboard nbextension'); - // }); - // test('Does not provide code lens for Python file launching tensorboard nbextension', async () => { - // const document = new MockDocument('a=1\n%tensorboard --logdir logs/fit', 'foo.py', async () => true); - // const codeLens = codeLensProvider.provideCodeLenses(document); - // assert.ok(codeLens.length === 0, 'Provided code lens for Python file loading tensorboard nbextension'); - // }); }); diff --git a/src/test/tensorBoard/tensorBoardImportCodeLensProvider.unit.test.ts b/src/test/tensorBoard/tensorBoardImportCodeLensProvider.unit.test.ts index 9b691c9af17c..07bcce035a7c 100644 --- a/src/test/tensorBoard/tensorBoardImportCodeLensProvider.unit.test.ts +++ b/src/test/tensorBoard/tensorBoardImportCodeLensProvider.unit.test.ts @@ -1,58 +1,72 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +import * as sinon from 'sinon'; import { assert } from 'chai'; import { CancellationTokenSource } from 'vscode'; +import { instance, mock } from 'ts-mockito'; import { TensorBoardImportCodeLensProvider } from '../../client/tensorBoard/tensorBoardImportCodeLensProvider'; import { MockDocument } from '../mocks/mockDocument'; +import { TensorboardExperiment } from '../../client/tensorBoard/tensorboarExperiment'; -suite('TensorBoard import code lens provider', () => { - let codeLensProvider: TensorBoardImportCodeLensProvider; - let cancelTokenSource: CancellationTokenSource; +[true, false].forEach((tbExtensionInstalled) => { + suite(`Tensorboard Extension is ${tbExtensionInstalled ? 'installed' : 'not installed'}`, () => { + suite.only('TensorBoard import code lens provider', () => { + let experiment: TensorboardExperiment; + let codeLensProvider: TensorBoardImportCodeLensProvider; + let cancelTokenSource: CancellationTokenSource; - setup(() => { - codeLensProvider = new TensorBoardImportCodeLensProvider([]); - cancelTokenSource = new CancellationTokenSource(); - }); - teardown(() => { - cancelTokenSource.dispose(); - }); - [ - 'import tensorboard', - 'import foo, tensorboard', - 'import foo, tensorboard, bar', - 'import tensorboardX', - 'import tensorboardX, bar', - 'import torch.profiler', - 'import foo, torch.profiler', - 'from torch.utils import tensorboard', - 'from torch.utils import foo, tensorboard', - 'import torch.utils.tensorboard, foo', - 'from torch import profiler', - ].forEach((importStatement) => { - test(`Provides code lens for Python files containing ${importStatement}`, () => { - const document = new MockDocument(importStatement, 'foo.py', async () => true); - const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); - assert.ok(codeLens.length > 0, `Failed to provide code lens for file containing ${importStatement} import`); - }); - test(`Provides code lens for Python ipynbs containing ${importStatement}`, () => { - const document = new MockDocument(importStatement, 'foo.ipynb', async () => true); - const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); - assert.ok( - codeLens.length > 0, - `Failed to provide code lens for ipynb containing ${importStatement} import`, - ); + setup(() => { + sinon.stub(TensorboardExperiment, 'isTensorboardExtensionInstalled').returns(tbExtensionInstalled); + experiment = mock(); + codeLensProvider = new TensorBoardImportCodeLensProvider([], instance(experiment)); + cancelTokenSource = new CancellationTokenSource(); + }); + teardown(() => { + sinon.restore(); + cancelTokenSource.dispose(); + }); + [ + 'import tensorboard', + 'import foo, tensorboard', + 'import foo, tensorboard, bar', + 'import tensorboardX', + 'import tensorboardX, bar', + 'import torch.profiler', + 'import foo, torch.profiler', + 'from torch.utils import tensorboard', + 'from torch.utils import foo, tensorboard', + 'import torch.utils.tensorboard, foo', + 'from torch import profiler', + ].forEach((importStatement) => { + test(`Provides code lens for Python files containing ${importStatement}`, () => { + const document = new MockDocument(importStatement, 'foo.py', async () => true); + const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); + assert.ok( + codeLens.length > 0, + `Failed to provide code lens for file containing ${importStatement} import`, + ); + }); + test(`Provides code lens for Python ipynbs containing ${importStatement}`, () => { + const document = new MockDocument(importStatement, 'foo.ipynb', async () => true); + const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); + assert.ok( + codeLens.length > 0, + `Failed to provide code lens for ipynb containing ${importStatement} import`, + ); + }); + test('Fails when cancellation is signaled', () => { + const document = new MockDocument(importStatement, 'foo.py', async () => true); + cancelTokenSource.cancel(); + const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); + assert.ok(codeLens.length === 0, 'Provided codelens even after cancellation was requested'); + }); + }); + test('Does not provide code lens if no matching import', () => { + const document = new MockDocument('import foo', 'foo.ipynb', async () => true); + const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); + assert.ok(codeLens.length === 0, 'Provided code lens for file without tensorboard import'); + }); }); - test('Fails when cancellation is signaled', () => { - const document = new MockDocument(importStatement, 'foo.py', async () => true); - cancelTokenSource.cancel(); - const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); - assert.ok(codeLens.length === 0, 'Provided codelens even after cancellation was requested'); - }); - }); - test('Does not provide code lens if no matching import', () => { - const document = new MockDocument('import foo', 'foo.ipynb', async () => true); - const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token); - assert.ok(codeLens.length === 0, 'Provided code lens for file without tensorboard import'); }); }); diff --git a/src/test/tensorBoard/tensorBoardPrompt.unit.test.ts b/src/test/tensorBoard/tensorBoardPrompt.unit.test.ts index 6f096e560d70..d94b0d6c5f23 100644 --- a/src/test/tensorBoard/tensorBoardPrompt.unit.test.ts +++ b/src/test/tensorBoard/tensorBoardPrompt.unit.test.ts @@ -7,7 +7,7 @@ import { Common } from '../../client/common/utils/localize'; import { TensorBoardEntrypointTrigger } from '../../client/tensorBoard/constants'; import { TensorBoardPrompt } from '../../client/tensorBoard/tensorBoardPrompt'; -suite('TensorBoard prompt', () => { +suite.only('TensorBoard prompt', () => { let applicationShell: ApplicationShell; let commandManager: CommandManager; let persistentState: PersistentState; diff --git a/src/test/tensorBoard/tensorBoardUsageTracker.unit.test.ts b/src/test/tensorBoard/tensorBoardUsageTracker.unit.test.ts index ff187dd2afc1..54771ab4b6b6 100644 --- a/src/test/tensorBoard/tensorBoardUsageTracker.unit.test.ts +++ b/src/test/tensorBoard/tensorBoardUsageTracker.unit.test.ts @@ -1,98 +1,117 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; -import { anything, reset, when } from 'ts-mockito'; +import { anything, instance, mock, reset, when } from 'ts-mockito'; import { TensorBoardUsageTracker } from '../../client/tensorBoard/tensorBoardUsageTracker'; import { TensorBoardPrompt } from '../../client/tensorBoard/tensorBoardPrompt'; import { MockDocumentManager } from '../mocks/mockDocumentManager'; import { createTensorBoardPromptWithMocks } from './helpers'; import { mockedVSCodeNamespaces } from '../vscode-mock'; +import { TensorboardExperiment } from '../../client/tensorBoard/tensorboarExperiment'; -suite('TensorBoard usage tracker', () => { - let documentManager: MockDocumentManager; - let tensorBoardImportTracker: TensorBoardUsageTracker; - let prompt: TensorBoardPrompt; - let showNativeTensorBoardPrompt: sinon.SinonSpy; +[true, false].forEach((tbExtensionInstalled) => { + suite(`Tensorboard Extension is ${tbExtensionInstalled ? 'installed' : 'not installed'}`, () => { + suite.only('TensorBoard usage tracker', () => { + let experiment: TensorboardExperiment; + let documentManager: MockDocumentManager; + let tensorBoardImportTracker: TensorBoardUsageTracker; + let prompt: TensorBoardPrompt; + let showNativeTensorBoardPrompt: sinon.SinonSpy; - suiteSetup(() => { - reset(mockedVSCodeNamespaces.extensions); - when(mockedVSCodeNamespaces.extensions?.getExtension(anything())).thenReturn(undefined); - }); - suiteTeardown(() => reset(mockedVSCodeNamespaces.extensions)); - setup(() => { - documentManager = new MockDocumentManager(); - prompt = createTensorBoardPromptWithMocks(); - showNativeTensorBoardPrompt = sinon.spy(prompt, 'showNativeTensorBoardPrompt'); - tensorBoardImportTracker = new TensorBoardUsageTracker(documentManager, [], prompt); - }); + suiteSetup(() => { + reset(mockedVSCodeNamespaces.extensions); + when(mockedVSCodeNamespaces.extensions?.getExtension(anything())).thenReturn(undefined); + }); + suiteTeardown(() => reset(mockedVSCodeNamespaces.extensions)); + setup(() => { + sinon.stub(TensorboardExperiment, 'isTensorboardExtensionInstalled').returns(tbExtensionInstalled); + experiment = mock(); + documentManager = new MockDocumentManager(); + prompt = createTensorBoardPromptWithMocks(); + showNativeTensorBoardPrompt = sinon.spy(prompt, 'showNativeTensorBoardPrompt'); + tensorBoardImportTracker = new TensorBoardUsageTracker( + documentManager, + [], + prompt, + instance(experiment), + ); + }); - test('Simple tensorboard import in Python file', async () => { - const document = documentManager.addDocument('import tensorboard', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('Simple tensorboardX import in Python file', async () => { - const document = documentManager.addDocument('import tensorboardX', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('Simple tensorboard import in Python ipynb', async () => { - const document = documentManager.addDocument('import tensorboard', 'foo.ipynb'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('`from x.y.tensorboard import z` import', async () => { - const document = documentManager.addDocument('from torch.utils.tensorboard import SummaryWriter', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('`from x.y import tensorboard` import', async () => { - const document = documentManager.addDocument('from torch.utils import tensorboard', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('`from tensorboardX import x` import', async () => { - const document = documentManager.addDocument('from tensorboardX import SummaryWriter', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('`import x, y` import', async () => { - const document = documentManager.addDocument('import tensorboard, tensorflow', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('`import pkg as _` import', async () => { - const document = documentManager.addDocument('import tensorboard as tb', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('Show prompt on changed text editor', async () => { - await tensorBoardImportTracker.activate(); - const document = documentManager.addDocument('import tensorboard as tb', 'foo.py'); - await documentManager.showTextDocument(document); - assert.ok(showNativeTensorBoardPrompt.calledOnce); - }); - test('Do not show prompt if no tensorboard import', async () => { - const document = documentManager.addDocument('import tensorflow as tf\nfrom torch.utils import foo', 'foo.py'); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.notCalled); - }); - test('Do not show prompt if language is not Python', async () => { - const document = documentManager.addDocument( - 'import tensorflow as tf\nfrom torch.utils import foo', - 'foo.cpp', - 'cpp', - ); - await documentManager.showTextDocument(document); - await tensorBoardImportTracker.activate(); - assert.ok(showNativeTensorBoardPrompt.notCalled); + test('Simple tensorboard import in Python file', async () => { + const document = documentManager.addDocument('import tensorboard', 'foo.py'); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('Simple tensorboardX import in Python file', async () => { + const document = documentManager.addDocument('import tensorboardX', 'foo.py'); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('Simple tensorboard import in Python ipynb', async () => { + const document = documentManager.addDocument('import tensorboard', 'foo.ipynb'); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('`from x.y.tensorboard import z` import', async () => { + const document = documentManager.addDocument( + 'from torch.utils.tensorboard import SummaryWriter', + 'foo.py', + ); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('`from x.y import tensorboard` import', async () => { + const document = documentManager.addDocument('from torch.utils import tensorboard', 'foo.py'); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('`from tensorboardX import x` import', async () => { + const document = documentManager.addDocument('from tensorboardX import SummaryWriter', 'foo.py'); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('`import x, y` import', async () => { + const document = documentManager.addDocument('import tensorboard, tensorflow', 'foo.py'); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('`import pkg as _` import', async () => { + const document = documentManager.addDocument('import tensorboard as tb', 'foo.py'); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('Show prompt on changed text editor', async () => { + await tensorBoardImportTracker.activate(); + const document = documentManager.addDocument('import tensorboard as tb', 'foo.py'); + await documentManager.showTextDocument(document); + assert.ok(showNativeTensorBoardPrompt.calledOnce); + }); + test('Do not show prompt if no tensorboard import', async () => { + const document = documentManager.addDocument( + 'import tensorflow as tf\nfrom torch.utils import foo', + 'foo.py', + ); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.notCalled); + }); + test('Do not show prompt if language is not Python', async () => { + const document = documentManager.addDocument( + 'import tensorflow as tf\nfrom torch.utils import foo', + 'foo.cpp', + 'cpp', + ); + await documentManager.showTextDocument(document); + await tensorBoardImportTracker.activate(); + assert.ok(showNativeTensorBoardPrompt.notCalled); + }); + }); }); });