Skip to content

Commit 552e278

Browse files
DonJayamanneDonJayamanne
committed
Tests
1 parent b123bf9 commit 552e278

File tree

5 files changed

+281
-77
lines changed

5 files changed

+281
-77
lines changed

‎src/client/datascience/interactive-ipynb/trustCommandHandler.ts

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@ import { inject, injectable } from 'inversify';
77
import { Uri } from 'vscode';
88
import { IExtensionSingleActivationService } from '../../activation/types';
99
import { IApplicationShell, ICommandManager } from '../../common/application/types';
10-
import { IDisposableRegistry } from '../../common/types';
10+
import { ContextKey } from '../../common/contextKey';
11+
import { EnableTrustedNotebooks } from '../../common/experiments/groups';
12+
import '../../common/extensions';
13+
import { IDisposableRegistry, IExperimentService } from '../../common/types';
1114
import { swallowExceptions } from '../../common/utils/decorators';
1215
import { DataScience } from '../../common/utils/localize';
1316
import { noop } from '../../common/utils/misc';
@@ -23,9 +26,18 @@ export class TrustCommandHandler implements IExtensionSingleActivationService {
2326
@inject(INotebookStorageProvider) private readonly storageProvider: INotebookStorageProvider,
2427
@inject(ICommandManager) private readonly commandManager: ICommandManager,
2528
@inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
26-
@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry
29+
@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
30+
@inject(IExperimentService) private readonly experiments: IExperimentService
2731
) {}
2832
public async activate(): Promise<void> {
33+
this.activateInBackground().ignoreErrors();
34+
}
35+
public async activateInBackground(): Promise<void> {
36+
if (!(await this.experiments.inExperiment(EnableTrustedNotebooks.experiment))) {
37+
return;
38+
}
39+
const context = new ContextKey('', this.commandManager);
40+
context.set(true).ignoreErrors();
2941
this.disposables.push(this.commandManager.registerCommand(Commands.TrustNotebook, this.onTrustNotebook, this));
3042
this.disposables.push(this.commandManager.registerCommand(Commands.TrustedNotebook, noop));
3143
}
@@ -35,15 +47,16 @@ export class TrustCommandHandler implements IExtensionSingleActivationService {
3547
if (!uri) {
3648
return;
3749
}
50+
3851
const model = await this.storageProvider.get(uri);
3952
if (model.isTrusted) {
4053
return;
4154
}
4255

43-
const prompts = [DataScience.trustNotebook(), DataScience.doNotTrustNotebook()];
4456
const selection = await this.applicationShell.showErrorMessage(
4557
DataScience.launchNotebookTrustPrompt(),
46-
...prompts
58+
DataScience.trustNotebook(),
59+
DataScience.doNotTrustNotebook()
4760
);
4861
if (selection !== DataScience.trustNotebook() || model.isTrusted) {
4962
return;

‎src/test/datascience/interactive-common/trustCommandHandler.unit.test.ts

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
'use strict';
4+
5+
import * as fakeTimers from '@sinonjs/fake-timers';
6+
import { assert } from 'chai';
7+
import * as sinon from 'sinon';
8+
import { anything, instance, mock, verify, when } from 'ts-mockito';
9+
import { Uri } from 'vscode';
10+
import { IExtensionSingleActivationService } from '../../../client/activation/types';
11+
import { IApplicationShell, ICommandManager } from '../../../client/common/application/types';
12+
import { ContextKey } from '../../../client/common/contextKey';
13+
import { EnableTrustedNotebooks } from '../../../client/common/experiments/groups';
14+
import { IDisposable, IExperimentService } from '../../../client/common/types';
15+
import { DataScience } from '../../../client/common/utils/localize';
16+
import { Commands } from '../../../client/datascience/constants';
17+
import { INotebookStorageProvider } from '../../../client/datascience/interactive-ipynb/notebookStorageProvider';
18+
import { TrustCommandHandler } from '../../../client/datascience/interactive-ipynb/trustCommandHandler';
19+
import { TrustService } from '../../../client/datascience/interactive-ipynb/trustService';
20+
import { INotebookEditorProvider, INotebookModel, ITrustService } from '../../../client/datascience/types';
21+
import { noop } from '../../core';
22+
import { createNotebookModel, disposeAllDisposables } from '../notebook/helper';
23+
24+
// tslint:disable: no-any
25+
26+
suite('DataScience - Trust Command Handler', () => {
27+
let trustCommandHandler: IExtensionSingleActivationService;
28+
let trustService: ITrustService;
29+
let editorProvider: INotebookEditorProvider;
30+
let storageProvider: INotebookStorageProvider;
31+
let commandManager: ICommandManager;
32+
let applicationShell: IApplicationShell;
33+
let disposables: IDisposable[] = [];
34+
let clock: fakeTimers.InstalledClock;
35+
let contextKeySet: sinon.SinonStub<[boolean], Promise<void>>;
36+
let experiments: IExperimentService;
37+
let model: INotebookModel;
38+
let trustNotebookCommandCallback: (uri: Uri) => Promise<void>;
39+
setup(() => {
40+
trustService = mock<TrustService>();
41+
editorProvider = mock<INotebookEditorProvider>();
42+
storageProvider = mock<INotebookStorageProvider>();
43+
commandManager = mock<ICommandManager>();
44+
applicationShell = mock<IApplicationShell>();
45+
model = createNotebookModel(false, Uri.file('a'));
46+
when(storageProvider.get(anything())).thenResolve(model);
47+
disposables = [];
48+
49+
experiments = mock<IExperimentService>();
50+
51+
when(trustService.trustNotebook(anything(), anything())).thenResolve();
52+
when(experiments.inExperiment(anything())).thenCall((exp) =>
53+
Promise.resolve(exp === EnableTrustedNotebooks.experiment)
54+
);
55+
when(commandManager.registerCommand(anything(), anything(), anything())).thenCall(() => ({ dispose: noop }));
56+
when(commandManager.registerCommand(Commands.TrustNotebook, anything(), anything())).thenCall((_, cb) => {
57+
trustNotebookCommandCallback = cb.bind(trustCommandHandler);
58+
return { dispose: noop };
59+
});
60+
61+
trustCommandHandler = new TrustCommandHandler(
62+
instance(trustService),
63+
instance(editorProvider),
64+
instance(storageProvider),
65+
instance(commandManager),
66+
instance(applicationShell),
67+
disposables,
68+
instance(experiments)
69+
);
70+
71+
clock = fakeTimers.install();
72+
73+
contextKeySet = sinon.stub(ContextKey.prototype, 'set');
74+
contextKeySet.resolves();
75+
});
76+
teardown(() => {
77+
sinon.restore();
78+
disposeAllDisposables(disposables);
79+
clock.uninstall();
80+
});
81+
82+
test('Context not set if not in experiment', async () => {
83+
when(experiments.inExperiment(anything())).thenResolve(false);
84+
85+
await trustCommandHandler.activate();
86+
await clock.runAllAsync();
87+
88+
assert.equal(contextKeySet.callCount, 0);
89+
});
90+
test('Context set if in experiment', async () => {
91+
when(experiments.inExperiment(anything())).thenCall((exp) =>
92+
Promise.resolve(exp === EnableTrustedNotebooks.experiment)
93+
);
94+
95+
await trustCommandHandler.activate();
96+
await clock.runAllAsync();
97+
98+
assert.equal(contextKeySet.callCount, 1);
99+
});
100+
test('Executing command will not update trust after dismissing the prompt', async () => {
101+
when(applicationShell.showErrorMessage(anything(), anything(), anything())).thenResolve(undefined as any);
102+
103+
await trustCommandHandler.activate();
104+
await clock.runAllAsync();
105+
await trustNotebookCommandCallback(Uri.file('a'));
106+
107+
verify(applicationShell.showErrorMessage(anything(), anything(), anything())).once();
108+
verify(trustService.trustNotebook(anything(), anything())).never();
109+
assert.isFalse(model.isTrusted);
110+
});
111+
test('Executing command will update trust', async () => {
112+
when(applicationShell.showErrorMessage(anything(), anything(), anything())).thenResolve(
113+
DataScience.trustNotebook() as any
114+
);
115+
116+
assert.isFalse(model.isTrusted);
117+
await trustCommandHandler.activate();
118+
await clock.runAllAsync();
119+
await trustNotebookCommandCallback(Uri.file('a'));
120+
121+
verify(applicationShell.showErrorMessage(anything(), anything(), anything())).once();
122+
verify(trustService.trustNotebook(anything(), anything())).once();
123+
assert.isTrue(model.isTrusted);
124+
});
125+
});

‎src/test/datascience/trustService.unit.test.ts renamed to ‎src/test/datascience/interactive-common/trustService.unit.test.ts

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,12 @@ import * as os from 'os';
88
import { anything, instance, mock, when } from 'ts-mockito';
99
import * as typemoq from 'typemoq';
1010
import { Uri } from 'vscode';
11-
import { ConfigurationService } from '../../client/common/configuration/service';
12-
import { ExperimentService } from '../../client/common/experiments/service';
13-
import { FileSystem } from '../../client/common/platform/fileSystem';
14-
import { IExtensionContext } from '../../client/common/types';
15-
import { DigestStorage } from '../../client/datascience/interactive-ipynb/digestStorage';
16-
import { TrustService } from '../../client/datascience/interactive-ipynb/trustService';
11+
import { ConfigurationService } from '../../../client/common/configuration/service';
12+
import { ExperimentService } from '../../../client/common/experiments/service';
13+
import { FileSystem } from '../../../client/common/platform/fileSystem';
14+
import { IExtensionContext } from '../../../client/common/types';
15+
import { DigestStorage } from '../../../client/datascience/interactive-ipynb/digestStorage';
16+
import { TrustService } from '../../../client/datascience/interactive-ipynb/trustService';
1717

1818
suite('DataScience - TrustService', () => {
1919
let trustService: TrustService;

‎src/test/datascience/notebook/helper.ts

Lines changed: 76 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,30 @@ import * as fs from 'fs-extra';
99
import * as path from 'path';
1010
import * as sinon from 'sinon';
1111
import * as tmp from 'tmp';
12-
import { commands, Uri } from 'vscode';
13-
import { NotebookCell } from '../../../../types/vscode-proposed';
12+
import { instance, mock } from 'ts-mockito';
13+
import { commands, TextDocument, Uri } from 'vscode';
14+
import { NotebookCell, NotebookDocument } from '../../../../types/vscode-proposed';
1415
import { CellDisplayOutput } from '../../../../typings/vscode-proposed';
1516
import { IApplicationEnvironment, IVSCodeNotebook } from '../../../client/common/application/types';
1617
import { MARKDOWN_LANGUAGE, PYTHON_LANGUAGE } from '../../../client/common/constants';
1718
import { IDisposable } from '../../../client/common/types';
1819
import { noop, swallowExceptions } from '../../../client/common/utils/misc';
19-
import { findMappedNotebookCellModel } from '../../../client/datascience/notebook/helpers/cellMappers';
20+
import { Identifiers } from '../../../client/datascience/constants';
21+
import { JupyterNotebookView } from '../../../client/datascience/notebook/constants';
22+
import {
23+
findMappedNotebookCellModel,
24+
mapVSCNotebookCellsToNotebookCellModels
25+
} from '../../../client/datascience/notebook/helpers/cellMappers';
26+
import { createVSCNotebookCellDataFromCell } from '../../../client/datascience/notebook/helpers/helpers';
2027
import { INotebookContentProvider } from '../../../client/datascience/notebook/types';
21-
import { ICell, INotebookEditorProvider, INotebookModel, INotebookProvider } from '../../../client/datascience/types';
28+
import { VSCodeNotebookModel } from '../../../client/datascience/notebookStorage/vscNotebookModel';
29+
import {
30+
CellState,
31+
ICell,
32+
INotebookEditorProvider,
33+
INotebookModel,
34+
INotebookProvider
35+
} from '../../../client/datascience/types';
2236
import { createEventHandler, waitForCondition } from '../../common';
2337
import { EXTENSION_ROOT_DIR_FOR_TESTS } from '../../constants';
2438
import { closeActiveWindows, initialize } from '../../initialize';
@@ -365,3 +379,61 @@ export async function saveActiveNotebook(disposables: IDisposable[]) {
365379

366380
await waitForCondition(async () => savedEvent.all.some((e) => e.kind === 'save'), 5_000, 'Not saved');
367381
}
382+
383+
export function createNotebookModel(trusted: boolean, uri: Uri, nb?: Partial<nbformat.INotebookContent>) {
384+
const nbJson: nbformat.INotebookContent = {
385+
cells: [],
386+
metadata: {
387+
orig_nbformat: 4
388+
},
389+
nbformat: 4,
390+
nbformat_minor: 4,
391+
...(nb || {})
392+
};
393+
394+
const cells = nbJson.cells.map((c, index) => {
395+
return {
396+
id: `NotebookImport#${index}`,
397+
file: Identifiers.EmptyFileName,
398+
line: 0,
399+
state: CellState.finished,
400+
data: c
401+
};
402+
});
403+
return new VSCodeNotebookModel(trusted, uri, JSON.parse(JSON.stringify(cells)));
404+
}
405+
406+
export function createNotebookDocument(
407+
model: INotebookModel,
408+
viewType: string = JupyterNotebookView
409+
): NotebookDocument {
410+
const doc: NotebookDocument = {
411+
cells: [],
412+
fileName: model.file.fsPath,
413+
isDirty: false,
414+
languages: [],
415+
uri: model.file,
416+
viewType,
417+
metadata: {
418+
cellEditable: model.isTrusted,
419+
cellHasExecutionOrder: true,
420+
cellRunnable: model.isTrusted,
421+
editable: model.isTrusted,
422+
runnable: model.isTrusted
423+
}
424+
};
425+
model.cells.forEach((cell, index) => {
426+
const vscCell = createVSCNotebookCellDataFromCell(model, cell)!;
427+
const vscDocumentCell: NotebookCell = {
428+
...vscCell,
429+
uri: model.file.with({ fragment: `cell${index}` }),
430+
notebook: doc,
431+
document: instance(mock<TextDocument>())
432+
};
433+
doc.cells.push(vscDocumentCell);
434+
});
435+
if (viewType === JupyterNotebookView) {
436+
mapVSCNotebookCellsToNotebookCellModels(doc, model);
437+
}
438+
return doc;
439+
}

0 commit comments

Comments
 (0)