Show Python: Launch TensorBoard command in command palette (#15382)

package.json

@@ -488,8 +488,7 @@
                 },
                 {
                     "command": "python.launchTensorBoard",
-                    "category": "Python",
-                    "when": "python.isInNativeTensorBoardExperiment"
+                    "category": "Python"
                 }
             ],
             "view/title": [

src/client/tensorBoard/tensorBoardSessionProvider.ts

@@ -5,11 +5,9 @@ import { inject, injectable } from 'inversify';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { IApplicationShell, ICommandManager, IWorkspaceService } from '../common/application/types';
 import { Commands } from '../common/constants';
-import { ContextKey } from '../common/contextKey';
-import { NativeTensorBoard } from '../common/experiments/groups';
 import { traceError, traceInfo } from '../common/logger';
 import { IProcessServiceFactory } from '../common/process/types';
-import { IDisposableRegistry, IExperimentService, IInstaller } from '../common/types';
+import { IDisposableRegistry, IInstaller } from '../common/types';
 import { TensorBoard } from '../common/utils/localize';
 import { IInterpreterService } from '../interpreter/contracts';
 import { sendTelemetryEvent } from '../telemetry';
@@ -26,34 +24,25 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
         @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
         @inject(ICommandManager) private readonly commandManager: ICommandManager,
         @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
-        @inject(IExperimentService) private readonly experimentService: IExperimentService,
         @inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory,
     ) {}
 
     public async activate(): Promise<void> {
-        this.activateInternal().ignoreErrors();
-    }
-
-    private async activateInternal() {
-        if (await this.experimentService.inExperiment(NativeTensorBoard.experiment)) {
-            this.disposables.push(
-                this.commandManager.registerCommand(
-                    Commands.LaunchTensorBoard,
-                    (
-                        entrypoint: TensorBoardEntrypoint = TensorBoardEntrypoint.palette,
-                        trigger: TensorBoardEntrypointTrigger = TensorBoardEntrypointTrigger.palette,
-                    ) => {
-                        sendTelemetryEvent(EventName.TENSORBOARD_SESSION_LAUNCH, undefined, {
-                            trigger,
-                            entrypoint,
-                        });
-                        return this.createNewSession();
-                    },
-                ),
-            );
-            const contextKey = new ContextKey('python.isInNativeTensorBoardExperiment', this.commandManager);
-            contextKey.set(true).ignoreErrors();
-        }
+        this.disposables.push(
+            this.commandManager.registerCommand(
+                Commands.LaunchTensorBoard,
+                (
+                    entrypoint: TensorBoardEntrypoint = TensorBoardEntrypoint.palette,
+                    trigger: TensorBoardEntrypointTrigger = TensorBoardEntrypointTrigger.palette,
+                ) => {
+                    sendTelemetryEvent(EventName.TENSORBOARD_SESSION_LAUNCH, undefined, {
+                        trigger,
+                        entrypoint,
+                    });
+                    return this.createNewSession();
+                },
+            ),
+        );
     }
 
     private async createNewSession(): Promise<TensorBoardSession | undefined> {

src/client/tensorBoard/tensorBoardUsageTracker.ts

@@ -7,7 +7,8 @@ import { TextEditor } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { IDocumentManager } from '../common/application/types';
 import { isTestExecution } from '../common/constants';
-import { IDisposableRegistry } from '../common/types';
+import { NativeTensorBoard } from '../common/experiments/groups';
+import { IDisposableRegistry, IExperimentService } from '../common/types';
 import { getDocumentLines } from '../telemetry/importTracker';
 import { TensorBoardEntrypointTrigger } from './constants';
 import { containsTensorBoardImport } from './helpers';
@@ -23,9 +24,13 @@ export class TensorBoardUsageTracker implements IExtensionSingleActivationServic
         @inject(IDocumentManager) private documentManager: IDocumentManager,
         @inject(IDisposableRegistry) private disposables: IDisposableRegistry,
         @inject(TensorBoardPrompt) private prompt: TensorBoardPrompt,
+        @inject(IExperimentService) private experimentService: IExperimentService,
     ) {}
 
     public async activate(): Promise<void> {
+        if (!(await this.experimentService.inExperiment(NativeTensorBoard.experiment))) {
+            return;
+        }
         if (testExecution) {
             await this.activateInternal();
         } else {

src/test/tensorBoard/tensorBoardSession.test.ts

@@ -7,7 +7,6 @@ import { TensorBoard } from '../../client/common/utils/localize';
 import { IServiceManager } from '../../client/ioc/types';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from '../../client/tensorBoard/constants';
 import { TensorBoardSession } from '../../client/tensorBoard/tensorBoardSession';
-import { TensorBoardSessionProvider } from '../../client/tensorBoard/tensorBoardSessionProvider';
 import { closeActiveWindows, initialize } from '../initialize';
 import * as ExperimentHelpers from '../../client/common/experiments/helpers';
 import { IInterpreterService } from '../../client/interpreter/contracts';
@@ -31,7 +30,6 @@ suite('TensorBoard session creation', async () => {
     let serviceManager: IServiceManager;
     let errorMessageStub: Sinon.SinonStub;
     let sandbox: Sinon.SinonSandbox;
-    let provider: TensorBoardSessionProvider;
     let applicationShell: IApplicationShell;
     let commandManager: ICommandManager;
 
@@ -59,9 +57,6 @@ suite('TensorBoard session creation', async () => {
         const interpreterService = serviceManager.get<IInterpreterService>(IInterpreterService);
         sandbox.stub(interpreterService, 'getActiveInterpreter').resolves(interpreter);
 
-        // Create tensorboard session provider
-        provider = serviceManager.get<TensorBoardSessionProvider>(TensorBoardSessionProvider);
-        await provider.activate();
         applicationShell = serviceManager.get<IApplicationShell>(IApplicationShell);
         errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
         commandManager = serviceManager.get<ICommandManager>(ICommandManager);

src/test/tensorBoard/tensorBoardUsageTracker.unit.test.ts

@@ -1,21 +1,32 @@
 import { assert } from 'chai';
 import * as sinon from 'sinon';
+import { mock, when, instance } from 'ts-mockito';
 import { TensorBoardUsageTracker } from '../../client/tensorBoard/tensorBoardUsageTracker';
 import { TensorBoardPrompt } from '../../client/tensorBoard/tensorBoardPrompt';
 import { MockDocumentManager } from '../startPage/mockDocumentManager';
 import { createTensorBoardPromptWithMocks } from './helpers';
+import { NativeTensorBoard } from '../../client/common/experiments/groups';
+import { ExperimentService } from '../../client/common/experiments/service';
 
 suite('TensorBoard usage tracker', () => {
     let documentManager: MockDocumentManager;
     let tensorBoardImportTracker: TensorBoardUsageTracker;
     let prompt: TensorBoardPrompt;
+    let experimentService: ExperimentService;
     let showNativeTensorBoardPrompt: sinon.SinonSpy;
 
     setup(() => {
         documentManager = new MockDocumentManager();
         prompt = createTensorBoardPromptWithMocks();
         showNativeTensorBoardPrompt = sinon.spy(prompt, 'showNativeTensorBoardPrompt');
-        tensorBoardImportTracker = new TensorBoardUsageTracker(documentManager, [], prompt);
+        experimentService = mock(ExperimentService);
+        when(experimentService.inExperiment(NativeTensorBoard.experiment)).thenResolve(true);
+        tensorBoardImportTracker = new TensorBoardUsageTracker(
+            documentManager,
+            [],
+            prompt,
+            instance(experimentService),
+        );
     });
 
     test('Simple tensorboard import in Python file', async () => {