Fail on cancellation signal when providing TensorBoard codelenses (microsoft#15332)

Author: joyceerhl

src/client/tensorBoard/nbextensionCodeLensProvider.ts

@@ -3,7 +3,7 @@
 
 import { inject, injectable } from 'inversify';
 import { once } from 'lodash';
-import { CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode';
+import { CancellationToken, CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { Commands, NotebookCellScheme, PYTHON_LANGUAGE } from '../common/constants';
 import { NativeTensorBoard } from '../common/experiments/groups';
@@ -46,7 +46,7 @@ export class TensorBoardNbextensionCodeLensProvider implements IExtensionSingleA
         }
     }
 
-    public provideCodeLenses(document: TextDocument): CodeLens[] {
+    public provideCodeLenses(document: TextDocument, cancelToken: CancellationToken): CodeLens[] {
         const command: Command = {
             title: TensorBoard.launchNativeTensorBoardSessionCodeLens(),
             command: Commands.LaunchTensorBoard,
@@ -56,6 +56,9 @@ export class TensorBoardNbextensionCodeLensProvider implements IExtensionSingleA
         };
         const codelenses: CodeLens[] = [];
         for (let index = 0; index < document.lineCount; index += 1) {
+            if (cancelToken.isCancellationRequested) {
+                return codelenses;
+            }
             const line = document.lineAt(index);
             if (containsNotebookExtension([line.text])) {
                 const range = new Range(new Position(line.lineNumber, 0), new Position(line.lineNumber, 1));

src/client/tensorBoard/tensorBoardImportCodeLensProvider.ts

@@ -3,7 +3,7 @@
 
 import { inject, injectable } from 'inversify';
 import { once } from 'lodash';
-import { CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode';
+import { CancellationToken, CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { Commands, PYTHON } from '../common/constants';
 import { NativeTensorBoard } from '../common/experiments/groups';
@@ -33,7 +33,7 @@ export class TensorBoardImportCodeLensProvider implements IExtensionSingleActiva
     }
 
     // eslint-disable-next-line class-methods-use-this
-    public provideCodeLenses(document: TextDocument): CodeLens[] {
+    public provideCodeLenses(document: TextDocument, cancelToken: CancellationToken): CodeLens[] {
         const command: Command = {
             title: TensorBoard.launchNativeTensorBoardSessionCodeLens(),
             command: Commands.LaunchTensorBoard,
@@ -43,6 +43,9 @@ export class TensorBoardImportCodeLensProvider implements IExtensionSingleActiva
         };
         const codelenses: CodeLens[] = [];
         for (let index = 0; index < document.lineCount; index += 1) {
+            if (cancelToken.isCancellationRequested) {
+                return codelenses;
+            }
             const line = document.lineAt(index);
             if (containsTensorBoardImport([line.text])) {
                 const range = new Range(new Position(line.lineNumber, 0), new Position(line.lineNumber, 1));

src/test/tensorBoard/nbextensionCodeLensProvider.unit.test.ts

@@ -3,6 +3,7 @@
 
 import { assert } from 'chai';
 import { mock } from 'ts-mockito';
+import { CancellationTokenSource } from 'vscode';
 import { ExperimentService } from '../../client/common/experiments/service';
 import { IExperimentService } from '../../client/common/types';
 import { TensorBoardNbextensionCodeLensProvider } from '../../client/tensorBoard/nbextensionCodeLensProvider';
@@ -11,32 +12,43 @@ import { MockDocument } from '../startPage/mockDocument';
 suite('TensorBoard nbextension code lens provider', () => {
     let experimentService: IExperimentService;
     let codeLensProvider: TensorBoardNbextensionCodeLensProvider;
+    let cancelTokenSource: CancellationTokenSource;
 
     setup(() => {
         experimentService = mock(ExperimentService);
         codeLensProvider = new TensorBoardNbextensionCodeLensProvider(experimentService, []);
+        cancelTokenSource = new CancellationTokenSource();
+    });
+    teardown(() => {
+        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 codeActions = codeLensProvider.provideCodeLenses(document);
-        assert.ok(codeActions.length > 0, 'Failed to provide code lens for file loading tensorboard nbextension');
+        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 codeActions = codeLensProvider.provideCodeLenses(document);
-        assert.ok(codeActions.length > 0, 'Failed to provide code lens for file loading tensorboard nbextension');
+        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 codeActions = codeLensProvider.provideCodeLenses(document);
-    //     assert.ok(codeActions.length === 0, 'Provided code lens for Python file loading tensorboard nbextension');
+    //     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 codeActions = codeLensProvider.provideCodeLenses(document);
-    //     assert.ok(codeActions.length === 0, 'Provided code lens for Python file loading tensorboard nbextension');
+    //     const codeLens = codeLensProvider.provideCodeLenses(document);
+    //     assert.ok(codeLens.length === 0, 'Provided code lens for Python file loading tensorboard nbextension');
     // });
 });

src/test/tensorBoard/tensorBoardImportCodeLensProvider.unit.test.ts

@@ -3,6 +3,7 @@
 
 import { assert } from 'chai';
 import { mock } from 'ts-mockito';
+import { CancellationTokenSource } from 'vscode';
 import { ExperimentService } from '../../client/common/experiments/service';
 import { IExperimentService } from '../../client/common/types';
 import { TensorBoardImportCodeLensProvider } from '../../client/tensorBoard/tensorBoardImportCodeLensProvider';
@@ -11,26 +12,37 @@ import { MockDocument } from '../startPage/mockDocument';
 suite('TensorBoard import code lens provider', () => {
     let experimentService: IExperimentService;
     let codeLensProvider: TensorBoardImportCodeLensProvider;
+    let cancelTokenSource: CancellationTokenSource;
 
     setup(() => {
         experimentService = mock(ExperimentService);
         codeLensProvider = new TensorBoardImportCodeLensProvider(experimentService, []);
+        cancelTokenSource = new CancellationTokenSource();
+    });
+    teardown(() => {
+        cancelTokenSource.dispose();
     });
     ['tensorboard', 'tensorboardX'].forEach((name) => {
         test(`Provides code lens for Python files importing ${name}`, () => {
             const document = new MockDocument(`import ${name}`, 'foo.py', async () => true);
-            const codeActions = codeLensProvider.provideCodeLenses(document);
-            assert.ok(codeActions.length > 0, `Failed to provide code lens for file containing ${name} import`);
+            const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token);
+            assert.ok(codeLens.length > 0, `Failed to provide code lens for file containing ${name} import`);
         });
         test(`Provides code lens for Python ipynbs importing ${name}`, () => {
             const document = new MockDocument(`import ${name}`, 'foo.ipynb', async () => true);
-            const codeActions = codeLensProvider.provideCodeLenses(document);
-            assert.ok(codeActions.length > 0, `Failed to provide code lens for ipynb containing ${name} import`);
+            const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token);
+            assert.ok(codeLens.length > 0, `Failed to provide code lens for ipynb containing ${name} import`);
+        });
+        test('Fails when cancellation is signaled', () => {
+            const document = new MockDocument(`import ${name}`, '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 codeActions = codeLensProvider.provideCodeLenses(document);
-        assert.ok(codeActions.length === 0, 'Provided code lens for file without tensorboard import');
+        const codeLens = codeLensProvider.provideCodeLenses(document, cancelTokenSource.token);
+        assert.ok(codeLens.length === 0, 'Provided code lens for file without tensorboard import');
     });
 });