Skip to content

Commit c655b55

Browse files
joyceerhljoyceerhl
authored
Better tests for trusted notebooks (microsoft#13158)
1 parent 7d807a8 commit c655b55

3 files changed

Lines changed: 196 additions & 48 deletions

File tree

‎news/3 Code Health/12966.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add better tests for trusted notebooks in the classic notebook editor.

‎src/client/datascience/notebookStorage/nativeEditorStorage.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,7 @@ export class NativeEditorStorage implements INotebookStorage {
189189
private async clearHotExit(file: Uri, backupId?: string): Promise<void> {
190190
const key = backupId || this.getStaticStorageKey(file);
191191
const filePath = this.getHashedFileName(key);
192-
await this.writeToStorage(filePath, undefined);
192+
await this.writeToStorage(filePath);
193193
}
194194

195195
private async writeToStorage(filePath: string, contents?: string, cancelToken?: CancellationToken): Promise<void> {

‎src/test/datascience/trustedNotebooks.functional.test.tsx‎

Lines changed: 194 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,35 @@
11
// Copyright (c) Microsoft Corporation. All rights reserved.
22
// Licensed under the MIT License.
33
'use strict';
4-
import { assert, use } from 'chai';
4+
import { assert, expect, use } from 'chai';
55
import * as chaiAsPromised from 'chai-as-promised';
66
import { ReactWrapper } from 'enzyme';
77
import * as fs from 'fs-extra';
88
import { Disposable } from 'vscode';
99
import { EnableTrustedNotebooks } from '../../client/common/experiments/groups';
1010
import { noop } from '../../client/common/utils/misc';
1111
import { InteractiveWindowMessages } from '../../client/datascience/interactive-common/interactiveWindowTypes';
12-
import { INotebookEditorProvider } from '../../client/datascience/types';
12+
import { INotebookEditor, INotebookEditorProvider, ITrustService } from '../../client/datascience/types';
13+
import { CommonActionType } from '../../datascience-ui/interactive-common/redux/reducers/types';
14+
import { TrustMessage } from '../../datascience-ui/interactive-common/trustMessage';
1315
import { NativeCell } from '../../datascience-ui/native-editor/nativeCell';
1416
import { NativeEditor } from '../../datascience-ui/native-editor/nativeEditor';
1517
import { IKeyboardEvent } from '../../datascience-ui/react-common/event';
1618
import { ImageButton } from '../../datascience-ui/react-common/imageButton';
1719
import { MonacoEditor } from '../../datascience-ui/react-common/monacoEditor';
1820
import { createTemporaryFile } from '../utils/fs';
1921
import { DataScienceIocContainer } from './dataScienceIocContainer';
20-
import { WaitForMessageOptions } from './mountedWebView';
21-
import { openEditor } from './nativeEditorTestHelpers';
22+
import { IMountedWebView, WaitForMessageOptions } from './mountedWebView';
23+
import { closeNotebook, openEditor } from './nativeEditorTestHelpers';
2224
import {
2325
addMockData,
2426
enterEditorKey,
2527
findButton,
2628
getNativeFocusedEditor,
2729
getOutputCell,
2830
isCellFocused,
31+
isCellMarkdown,
32+
isCellSelected,
2933
typeCode,
3034
verifyHtmlOnCell
3135
} from './testHelpers';
@@ -40,8 +44,9 @@ function waitForMessage(ioc: DataScienceIocContainer, message: string, options?:
4044
.waitForMessage(message, options);
4145
}
4246
// tslint:disable:no-any no-multiline-string
43-
suite('Untrusted notebooks', () => {
47+
suite('Notebook trust', () => {
4448
let wrapper: ReactWrapper<any, Readonly<{}>, React.Component>;
49+
let ne: { editor: INotebookEditor; mount: IMountedWebView };
4550
const disposables: Disposable[] = [];
4651
let ioc: DataScienceIocContainer;
4752
const baseFile = `
@@ -217,22 +222,31 @@ suite('Untrusted notebooks', () => {
217222
wrapper.update();
218223
}
219224

220-
async function setupFunction(this: Mocha.Context, fileContents?: any) {
225+
async function setupFunction(this: Mocha.Context) {
221226
addMockData(ioc, 'b=2\nb', 2);
222227
addMockData(ioc, 'c=3\nc', 3);
223228
// Use a real file so we can save notebook to a file.
224229
// This is used in some tests (saving).
225230
notebookFile = await createTemporaryFile('.ipynb');
226-
await fs.writeFile(notebookFile.filePath, fileContents ? fileContents : baseFile);
227-
const ne = await openEditor(ioc, fileContents ? fileContents : baseFile, notebookFile.filePath);
231+
await fs.writeFile(notebookFile.filePath, baseFile);
232+
ne = await openEditor(ioc, baseFile, notebookFile.filePath);
228233
wrapper = ne.mount.wrapper;
229234
}
235+
230236
function clickCell(cellIndex: number) {
231237
wrapper.update();
232238
wrapper.find(NativeCell).at(cellIndex).simulate('click');
233239
wrapper.update();
234240
}
235241

242+
async function focusCell(targetCellIndex: number) {
243+
const update = waitForMessage(ioc, InteractiveWindowMessages.FocusedCellEditor);
244+
clickCell(targetCellIndex);
245+
simulateKeyPressOnCell(targetCellIndex, { code: 'Enter', editorInfo: undefined });
246+
await update;
247+
assert.ok(isCellFocused(wrapper, 'NativeCell', targetCellIndex));
248+
}
249+
236250
teardown(async () => {
237251
for (const disposable of disposables) {
238252
if (!disposable) {
@@ -258,47 +272,180 @@ suite('Untrusted notebooks', () => {
258272
await setupFunction.call(this);
259273
});
260274

261-
test('Outputs are not rendered', () => {
262-
// No outputs should have rendered
263-
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>1</span>', 0));
264-
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>2</span>', 1));
265-
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>3</span>', 2));
266-
});
267-
test('Cannot insert cell into notebook', async () => {
268-
// Click insert cell button
269-
const insertCellButton = findButton(wrapper, NativeEditor, 5);
270-
insertCellButton?.simulate('click');
275+
suite('Open an untrusted notebook', async () => {
276+
test('Outputs are not rendered', () => {
277+
// No outputs should have rendered
278+
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>1</span>', 0));
279+
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>2</span>', 1));
280+
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>3</span>', 2));
281+
});
282+
test('Cannot edit cell contents', async () => {
283+
await focusCell(0);
271284

272-
// No cell should have been added
273-
assert.equal(wrapper.find('NativeCell').length, 3, 'Cell added');
274-
});
275-
test('Cannot edit cell contents', async () => {
276-
// Set focus to a cell
277-
let update = waitForMessage(ioc, InteractiveWindowMessages.SelectedCell);
278-
clickCell(0);
279-
await update;
280-
update = waitForMessage(ioc, InteractiveWindowMessages.FocusedCellEditor);
281-
simulateKeyPressOnCell(0, { code: 'Enter' });
282-
await update;
283-
assert.ok(isCellFocused(wrapper, 'NativeCell', 0));
285+
// Try to type code
286+
const editorEnzyme = getNativeFocusedEditor(wrapper);
287+
typeCode(editorEnzyme, 'foo');
288+
const reactEditor = editorEnzyme!.instance() as MonacoEditor;
289+
const editor = reactEditor.state.editor;
290+
if (editor) {
291+
assert.notInclude(editor.getModel()!.getValue(), 'foo', 'Was able to edit cell in untrusted notebook');
292+
}
293+
});
284294

285-
// Then try to type code
286-
const editorEnzyme = getNativeFocusedEditor(wrapper);
287-
typeCode(editorEnzyme, 'foo');
288-
const reactEditor = editorEnzyme!.instance() as MonacoEditor;
289-
const editor = reactEditor.state.editor;
290-
if (editor) {
291-
assert.notInclude(editor.getModel()!.getValue(), 'foo', 'Was able to edit cell in untrusted notebook');
292-
}
295+
suite('Buttons are disabled', async () => {
296+
test('Cannot run cell', async () => {
297+
// Click run cell button
298+
const cell = getOutputCell(wrapper, 'NativeCell', 1);
299+
const imageButtons = cell!.find(ImageButton);
300+
const runButton = imageButtons.findWhere((w) => w.props().tooltip === 'Run cell');
301+
runButton.simulate('click');
302+
303+
// Ensure cell was not executed
304+
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', `2`, 1));
305+
});
306+
test('Cannot switch to markdown', async () => {
307+
// Click switch to markdown button
308+
const cell = getOutputCell(wrapper, 'NativeCell', 1);
309+
const imageButtons = cell!.find(ImageButton);
310+
const changeToMarkdown = imageButtons.findWhere((w) => w.props().tooltip === 'Change to markdown');
311+
changeToMarkdown.simulate('click');
312+
313+
// Ensure cell is still code cell
314+
assert.isFalse(isCellMarkdown(wrapper, 'NativeCell', 1));
315+
});
316+
test('Cannot insert cell into notebook', async () => {
317+
// Click insert cell button
318+
const insertCellButton = findButton(wrapper, NativeEditor, 5);
319+
insertCellButton?.simulate('click');
320+
321+
// No cell should have been added
322+
assert.equal(wrapper.find('NativeCell').length, 3, 'Cell added');
323+
});
324+
});
325+
326+
suite('Jupyter shortcuts for editing notebook are disabled', async () => {
327+
test('Ctrl+enter does not execute cell', async () => {
328+
const cellIndex = 1;
329+
await focusCell(cellIndex);
330+
331+
const promise = waitForMessage(ioc, InteractiveWindowMessages.ExecutionRendered, { timeoutMs: 5_000 });
332+
simulateKeyPressOnCell(cellIndex, { code: 'Enter', ctrlKey: true, editorInfo: undefined });
333+
334+
// Waiting for an execution rendered message should timeout
335+
await expect(promise).to.eventually.be.rejected;
336+
// No output should have been rendered
337+
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>2</span>', cellIndex));
338+
});
339+
test('Shift+enter does not execute cell or advance to next cell', async () => {
340+
const cellIndex = 1;
341+
await focusCell(cellIndex);
342+
343+
const promise = waitForMessage(ioc, InteractiveWindowMessages.ExecutionRendered, { timeoutMs: 5_000 });
344+
simulateKeyPressOnCell(cellIndex, { code: 'Enter', shiftKey: true, editorInfo: undefined });
345+
346+
// Waiting for an execution rendered message should timeout
347+
await expect(promise).to.eventually.be.rejected;
348+
// No output should have been rendered
349+
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>2</span>', cellIndex));
350+
// 3rd cell should be neither selected nor focused
351+
assert.isFalse(isCellSelected(wrapper, 'NativeCell', cellIndex + 1));
352+
assert.isFalse(isCellFocused(wrapper, 'NativeCell', cellIndex + 1));
353+
});
354+
test('Alt+enter does not execute cell or add a new cell below', async () => {
355+
assert.equal(wrapper.find('NativeCell').length, 3);
356+
const cellIndex = 1;
357+
await focusCell(cellIndex);
358+
359+
const promise = waitForMessage(ioc, InteractiveWindowMessages.ExecutionRendered, { timeoutMs: 5_000 });
360+
simulateKeyPressOnCell(1, { code: 'Enter', altKey: true, editorInfo: undefined });
361+
362+
// Waiting for an execution rendered message should timeout
363+
await expect(promise).to.eventually.be.rejected;
364+
// No output should have been rendered
365+
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', '<span>2</span>', cellIndex));
366+
// No cell should have been added
367+
assert.equal(wrapper.find('NativeCell').length, 3, 'Cell added');
368+
});
369+
test('"a" does not add a cell', async () => {
370+
assert.equal(wrapper.find('NativeCell').length, 3);
371+
const cellIndex = 0;
372+
await focusCell(cellIndex);
373+
374+
const addedCell = waitForMessage(ioc, CommonActionType.INSERT_ABOVE_AND_FOCUS_NEW_CELL, {
375+
timeoutMs: 5_000
376+
});
377+
const update = waitForMessage(ioc, InteractiveWindowMessages.SelectedCell, { timeoutMs: 5_000 });
378+
simulateKeyPressOnCell(cellIndex, { code: 'a' });
379+
380+
await expect(addedCell).to.eventually.be.rejected;
381+
await expect(update).to.eventually.be.rejected;
382+
// There should still be 3 cells.
383+
assert.equal(wrapper.find('NativeCell').length, 3);
384+
});
385+
test('"b" does not add a cell', async () => {
386+
assert.equal(wrapper.find('NativeCell').length, 3);
387+
const cellIndex = 0;
388+
await focusCell(cellIndex);
389+
390+
const addedCell = waitForMessage(ioc, CommonActionType.INSERT_BELOW_AND_FOCUS_NEW_CELL, {
391+
timeoutMs: 5_000
392+
});
393+
const update = waitForMessage(ioc, InteractiveWindowMessages.SelectedCell, { timeoutMs: 5_000 });
394+
simulateKeyPressOnCell(cellIndex, { code: 'b' });
395+
396+
await expect(addedCell).to.eventually.be.rejected;
397+
await expect(update).to.eventually.be.rejected;
398+
// There should still be 3 cells.
399+
assert.equal(wrapper.find('NativeCell').length, 3);
400+
});
401+
test('"d" does not delete a cell', async () => {
402+
assert.equal(wrapper.find('NativeCell').length, 3);
403+
const cellIndex = 2;
404+
await focusCell(cellIndex);
405+
406+
simulateKeyPressOnCell(cellIndex, { code: 'd' });
407+
simulateKeyPressOnCell(cellIndex, { code: 'd' });
408+
409+
// There should still be 3 cells.
410+
assert.equal(wrapper.find('NativeCell').length, 3);
411+
});
412+
test('"m" does not change a code cell to markdown', async () => {
413+
const cellIndex = 2;
414+
await focusCell(cellIndex);
415+
416+
const update = waitForMessage(ioc, CommonActionType.CHANGE_CELL_TYPE, {
417+
timeoutMs: 5_000
418+
});
419+
simulateKeyPressOnCell(cellIndex, { code: 'm' });
420+
421+
await expect(update).to.eventually.be.rejected;
422+
assert.isFalse(
423+
isCellMarkdown(wrapper, 'NativeCell', cellIndex),
424+
'Code cell in untrusted notebook was changed to markdown'
425+
);
426+
});
427+
});
293428
});
294-
test('Cannot run cell', async () => {
295-
// Click run cell button
296-
const cell = getOutputCell(wrapper, 'NativeCell', 1);
297-
const imageButtons = cell!.find(ImageButton);
298-
const runButton = imageButtons.findWhere((w) => w.props().tooltip === 'Run cell');
299-
runButton.simulate('click');
300-
301-
// Ensure cell was not executed
302-
assert.throws(() => verifyHtmlOnCell(wrapper, 'NativeCell', `2`, 1));
429+
430+
suite('Trust an untrusted notebook', async () => {
431+
test('Trust persists when closed and reopened', async () => {
432+
const before = wrapper.find(TrustMessage);
433+
assert.equal(before.text(), 'Not Trusted');
434+
435+
// Trust notebook
436+
const trustService = ioc.get<ITrustService>(ITrustService);
437+
await trustService.trustNotebook(ne.editor.model.file, ne.editor.model.getContent());
438+
439+
// Close
440+
await closeNotebook(ioc, ne.editor);
441+
442+
// Reopen
443+
const newNativeEditor = await openEditor(ioc, baseFile, notebookFile.filePath);
444+
const newWrapper = newNativeEditor.mount.wrapper;
445+
446+
// Verify notebook is now trusted
447+
const after = newWrapper.find(TrustMessage);
448+
assert.equal(after.text(), 'Trusted');
449+
});
303450
});
304451
});

0 commit comments

Comments
 (0)