// Copyright (c) Microsoft Corporation. All rights reserved.

// Licensed under the MIT License.

'use strict';

import { nbformat } from '@jupyterlab/coreutils';

import { assert } from 'chai';

import { ChildProcess } from 'child_process';

import * as fs from 'fs-extra';

import { injectable } from 'inversify';

import * as os from 'os';

import * as path from 'path';

import { SemVer } from 'semver';

import { Readable, Writable } from 'stream';

import { anything, instance, mock, when } from 'ts-mockito';

import * as TypeMoq from 'typemoq';

import * as uuid from 'uuid/v4';

import { Disposable, Uri } from 'vscode';

import { CancellationToken, CancellationTokenSource } from 'vscode-jsonrpc';

import { ApplicationShell } from '../../client/common/application/applicationShell';

import { IApplicationShell } from '../../client/common/application/types';

import { Cancellation, CancellationError } from '../../client/common/cancellation';

import { EXTENSION_ROOT_DIR } from '../../client/common/constants';

import { LocalZMQKernel } from '../../client/common/experiments/groups';

import { traceError, traceInfo } from '../../client/common/logger';

import { IFileSystem } from '../../client/common/platform/types';

import { IPythonExecutionFactory, IPythonExecutionService, Output } from '../../client/common/process/types';

import { Product } from '../../client/common/types';

import { createDeferred, waitForPromise } from '../../client/common/utils/async';

import { noop } from '../../client/common/utils/misc';

import { Architecture } from '../../client/common/utils/platform';

import { Identifiers } from '../../client/datascience/constants';

import { getMessageForLibrariesNotInstalled } from '../../client/datascience/jupyter/interpreter/jupyterInterpreterDependencyService';

import { JupyterExecutionFactory } from '../../client/datascience/jupyter/jupyterExecutionFactory';

import { JupyterKernelPromiseFailedError } from '../../client/datascience/jupyter/kernels/jupyterKernelPromiseFailedError';

import { HostJupyterNotebook } from '../../client/datascience/jupyter/liveshare/hostJupyterNotebook';

import {

CellState,

ICell,

IJupyterConnection,

IJupyterExecution,

IJupyterKernelSpec,

INotebook,

INotebookExecutionLogger,

INotebookExporter,

INotebookImporter,

INotebookProvider,

InterruptResult

} from '../../client/datascience/types';

import { IInterpreterService, IKnownSearchPathsForInterpreters } from '../../client/interpreter/contracts';

import { InterpreterType, PythonInterpreter } from '../../client/pythonEnvironments/discovery/types';

import { concatMultilineStringInput } from '../../datascience-ui/common';

import { generateTestState, ICellViewModel } from '../../datascience-ui/interactive-common/mainState';

import { sleep } from '../core';

import { DataScienceIocContainer } from './dataScienceIocContainer';

import { takeSnapshot, writeDiffSnapshot } from './helpers';

import { getIPConnectionInfo } from './jupyterHelpers';

import { SupportedCommands } from './mockJupyterManager';

import { MockPythonService } from './mockPythonService';

// tslint:disable:no-any no-multiline-string max-func-body-length no-console max-classes-per-file trailing-comma

suite('DataScience notebook tests', () => {

[false, true].forEach((useRawKernel) => {

suite(`${useRawKernel ? 'With Direct Kernel' : 'With Jupyter Server'}`, () => {

const disposables: Disposable[] = [];

let notebookProvider: INotebookProvider;

let pythonFactory: IPythonExecutionFactory;

let ioc: DataScienceIocContainer;

let modifiedConfig = false;

const baseUri = Uri.file('foo.py');

let snapshot: any;

// tslint:disable-next-line: no-function-expression

setup(async function () {

ioc = new DataScienceIocContainer();

if (ioc.shouldMockJupyter && useRawKernel) {

// tslint:disable-next-line: no-invalid-this

this.skip();

return;

} else {

ioc.setExperimentState(LocalZMQKernel.experiment, useRawKernel);

}

ioc.registerDataScienceTypes();

await ioc.activate();

notebookProvider = ioc.get<INotebookProvider>(INotebookProvider);

pythonFactory = ioc.get<IPythonExecutionFactory>(IPythonExecutionFactory);

});

suiteSetup(() => {

snapshot = takeSnapshot();

});

suiteTeardown(() => {

writeDiffSnapshot(snapshot, `Notebook ${useRawKernel}`);

});

teardown(async () => {

try {

if (modifiedConfig) {

traceInfo('Attempting to put jupyter default config back');

const procService = await createPythonService();

if (procService) {

await procService.exec(['-m', 'jupyter', 'notebook', '--generate-config', '-y'], {});

}

}

traceInfo('Shutting down after test.');

// tslint:disable-next-line:prefer-for-of

for (let i = 0; i < disposables.length; i += 1) {

const disposable = disposables[i];

if (disposable) {

const promise = disposable.dispose() as Promise<any>;

if (promise) {

await promise;

}

}

}

await ioc.dispose();

traceInfo('Shutdown after test complete.');

} catch (e) {

traceError(e);

}

if (process.env.PYTHONWARNINGS) {

delete process.env.PYTHONWARNINGS;

}

});

function escapePath(p: string) {

return p.replace(/\\/g, '\\\\');

}

function srcDirectory() {

return path.join(EXTENSION_ROOT_DIR, 'src', 'test', 'datascience');

}

function extractDataOutput(cell: ICell): any {

assert.equal(cell.data.cell_type, 'code', `Wrong type of cell returned`);

const codeCell = cell.data as nbformat.ICodeCell;

if (codeCell.outputs.length > 0) {

assert.equal(codeCell.outputs.length, 1, 'Cell length not correct');

const data = codeCell.outputs[0].data;

const error = codeCell.outputs[0].evalue;

if (error) {

assert.fail(`Unexpected error: ${error}`);

}

assert.ok(data, `No data object on the cell`);

if (data) {

// For linter

assert.ok(data.hasOwnProperty('text/plain'), `Cell mime type not correct`);

assert.ok((data as any)['text/plain'], `Cell mime type not correct`);

return (data as any)['text/plain'];

}

}

}

async function verifySimple(

notebook: INotebook | undefined,

code: string,

expectedValue: any,

pathVerify = false

): Promise<void> {

const cells = await notebook!.execute(code, path.join(srcDirectory(), 'foo.py'), 2, uuid());

assert.equal(cells.length, 1, `Wrong number of cells returned`);

const data = extractDataOutput(cells[0]);

if (pathVerify) {

// For a path comparison normalize output and add single quotes on expected value

const normalizedOutput = path.normalize(data).toUpperCase();

const normalizedTarget = `'${path.normalize(expectedValue).toUpperCase()}'`;

assert.equal(normalizedOutput, normalizedTarget, 'Cell path values does not match');

} else {

assert.equal(data, expectedValue, 'Cell value does not match');

}

}

async function verifyError(

notebook: INotebook | undefined,

code: string,

errorString: string

): Promise<void> {

const cells = await notebook!.execute(code, path.join(srcDirectory(), 'foo.py'), 2, uuid());

assert.equal(cells.length, 1, `Wrong number of cells returned`);

assert.equal(cells[0].data.cell_type, 'code', `Wrong type of cell returned`);

const cell = cells[0].data as nbformat.ICodeCell;

assert.equal(cell.outputs.length, 1, `Cell length not correct`);

const error = cell.outputs[0].evalue;

if (error) {

assert.ok(error, 'Error not found when expected');

assert.equal(error, errorString, 'Unexpected error found');

}

}

async function verifyCell(

notebook: INotebook | undefined,

index: number,

code: string,

mimeType: string,

cellType: string,

verifyValue: (data: any) => void

): Promise<void> {

// Verify results of an execute

const cells = await notebook!.execute(code, path.join(srcDirectory(), 'foo.py'), 2, uuid());

assert.equal(cells.length, 1, `${index}: Wrong number of cells returned`);

if (cellType === 'code') {

assert.equal(cells[0].data.cell_type, cellType, `${index}: Wrong type of cell returned`);

const cell = cells[0].data as nbformat.ICodeCell;

assert.ok(cell.outputs.length >= 1, `${index}: Cell length not correct`);

const error = cell.outputs[0].evalue;

if (error) {

assert.ok(false, `${index}: Unexpected error: ${error}`);

}

const data = cell.outputs[0].data;

const text = cell.outputs[0].text;

assert.ok(data || text, `${index}: No data object on the cell for ${code}`);

if (data) {

// For linter

assert.ok(

data.hasOwnProperty(mimeType) || data.hasOwnProperty('text/plain'),

`${index}: Cell mime type not correct for ${JSON.stringify(data)}`

);

const actualMimeType = data.hasOwnProperty(mimeType) ? mimeType : 'text/plain';

assert.ok((data as any)[actualMimeType], `${index}: Cell mime type not correct`);

verifyValue((data as any)[actualMimeType]);

}

if (text) {

verifyValue(text);

}

} else if (cellType === 'markdown') {

assert.equal(cells[0].data.cell_type, cellType, `${index}: Wrong type of cell returned`);

const cell = cells[0].data as nbformat.IMarkdownCell;

const outputSource = concatMultilineStringInput(cell.source);

verifyValue(outputSource);

} else if (cellType === 'error') {

const cell = cells[0].data as nbformat.ICodeCell;

assert.equal(cell.outputs.length, 1, `${index}: Cell length not correct`);

const error = cell.outputs[0].evalue;

assert.ok(error, 'Error not found when expected');

verifyValue(error);

}

}

function testMimeTypes(

types: {

markdownRegEx: string | undefined;

code: string;

mimeType: string;

result: any;

cellType: string;

verifyValue(data: any): void;

}[]

) {

runTest('MimeTypes', async () => {

// Prefill with the output (This is only necessary for mocking)

types.forEach((t) => {

addMockData(t.code, t.result, t.mimeType, t.cellType);

});

// Test all mime types together so we don't have to startup and shutdown between

// each

const server = await createNotebook();

if (server) {

for (let i = 0; i < types.length; i += 1) {

const markdownRegex = types[i].markdownRegEx ? types[i].markdownRegEx : '';

ioc.getSettings().datascience.markdownRegularExpression = markdownRegex!;

await verifyCell(

server,

i,

types[i].code,

types[i].mimeType,

types[i].cellType,

types[i].verifyValue

);

}

}

});

}

function runTest(

name: string,

func: (_this: Mocha.Context) => Promise<void>,

_notebookProc?: ChildProcess,

rebindFunc?: () => void

) {

test(name, async function () {

// Give tests a chance to rebind IOC services before we fetch jupyterExecution and processFactory

if (rebindFunc) {

rebindFunc();

}

console.log(`Starting test ${name} ...`);

// tslint:disable-next-line: no-invalid-this

return func(this);

});

}

async function createNotebookWithNonDefaultConfig(): Promise<INotebook | undefined> {

const newSettings = { ...ioc.getSettings().datascience, useDefaultConfig: false };

ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath, newSettings);

return createNotebook();

}

async function createNotebook(

uri?: string,

launchingFile?: string,

expectFailure?: boolean

): Promise<INotebook | undefined> {

// Catch exceptions. Throw a specific assertion if the promise fails

try {

if (uri) {

const newSettings = { ...ioc.getSettings().datascience, jupyterServerURI: uri };

ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath, newSettings);

}

const notebook = await notebookProvider.getOrCreateNotebook({

identity: Uri.parse(Identifiers.InteractiveWindowIdentity)

});

launchingFile = launchingFile || path.join(srcDirectory(), 'foo.py');

if (notebook) {

await notebook.setLaunchingFile(launchingFile);

}

return notebook;

} catch (exc) {

if (!expectFailure) {

assert.ok(false, `Expected server to be created, but got ${exc}`);

}

}

}

function addMockData(code: string, result: string | number, mimeType?: string, cellType?: string) {

if (ioc.mockJupyter) {

if (cellType && cellType === 'error') {

ioc.mockJupyter.addError(code, result.toString());

} else {

ioc.mockJupyter.addCell(code, result, mimeType);

}

}

}

function changeMockWorkingDirectory(workingDir: string) {

if (ioc.mockJupyter) {

ioc.mockJupyter.changeWorkingDirectory(workingDir);

}

}

function addInterruptableMockData(

code: string,

resultGenerator: (c: CancellationToken) => Promise<{ result: string; haveMore: boolean }>

) {

if (ioc.mockJupyter) {

ioc.mockJupyter.addContinuousOutputCell(code, resultGenerator);

}

}

async function createPythonService(

versionRequirement?: number

): Promise<IPythonExecutionService | undefined> {

if (!ioc.mockJupyter) {

const python = await ioc.getJupyterCapableInterpreter();

if (

python &&

python.version?.major &&

(!versionRequirement || python.version?.major > versionRequirement)

) {

return pythonFactory.createActivatedEnvironment({

resource: undefined,

interpreter: python,

allowEnvironmentFetchExceptions: true,

bypassCondaExecution: true

});

}

}

}

async function startRemoteServer(pythonService: IPythonExecutionService, args: string[]): Promise<string> {

const connectionFound = createDeferred();

const exeResult = pythonService.execObservable(args, {

throwOnStdErr: false

});

disposables.push(exeResult);

exeResult.out.subscribe(

(output: Output<string>) => {

traceInfo(`Remote server output: ${output.out}`);

const connectionURL = getIPConnectionInfo(output.out);

if (connectionURL) {

connectionFound.resolve(connectionURL);

}

},

(e) => {

traceInfo(`Remote server error: ${e}`);

connectionFound.reject(e);

}

);

traceInfo('Connecting to remote server');

const connString = await connectionFound.promise;

const uri = connString as string;

// Wait another 3 seconds to give notebook time to be ready. Not sure

// how else to know when it's okay to connect to. Mac on azure seems

// to connect too fast and then is unable to actually communicate.

await sleep(3000);

return uri;

}

runTest('Remote Self Certs', async (_this: Mocha.Context) => {

const pythonService = await createPythonService(2);

// Skip test for older python and raw kernel

if (pythonService && !useRawKernel) {

// We will only connect if we allow for self signed cert connections

ioc.forceDataScienceSettingsChanged({

allowUnauthorizedRemoteConnection: true,

jupyterLaunchTimeout: 60000

});

const pemFile = path.join(

EXTENSION_ROOT_DIR,

'src',

'test',

'datascience',

'serverConfigFiles',

'jcert.pem'

);

const keyFile = path.join(

EXTENSION_ROOT_DIR,

'src',

'test',

'datascience',

'serverConfigFiles',

'jkey.key'

);

const uri = await startRemoteServer(pythonService, [

'-m',

'jupyter',

'notebook',

'--NotebookApp.open_browser=False',

'--NotebookApp.ip=*',

'--NotebookApp.port=9999',

`--certfile=${pemFile}`,

`--keyfile=${keyFile}`

]);

traceInfo('Waiting for notebook');

// We have a connection string here, so try to connect jupyterExecution to the notebook server

const notebook = await createNotebook(uri);

if (!notebook) {

assert.fail(`Failed to connect to remote self cert server on ${uri}`);

} else {

await verifySimple(notebook, `a=1${os.EOL}a`, 1);

}

} else {

traceInfo('Remote Self Cert is not supported on 2.7');

_this.skip();

}

});

// Connect to a server that doesn't have a token or password, customers use this and we regressed it once

runTest(

'Remote No Auth',

async () => {

const pythonService = await createPythonService();

if (pythonService) {

const configFile = path.join(

EXTENSION_ROOT_DIR,

'src',

'test',

'datascience',

'serverConfigFiles',

'remoteNoAuth.py'

);

const uri = await startRemoteServer(pythonService, [

'-m',

'jupyter',

'notebook',

`--config=${configFile}`

]);

// We have a connection string here, so try to connect jupyterExecution to the notebook server

const notebook = await createNotebook(uri);

if (!notebook) {

assert.fail('Failed to connect to remote password server');

} else {

await verifySimple(notebook, `a=1${os.EOL}a`, 1);

}

}

},

undefined,

() => {

const dummyDisposable = {

dispose: () => {

return;

}

};

const appShell = TypeMoq.Mock.ofType<IApplicationShell>();

appShell

.setup((a) => a.showErrorMessage(TypeMoq.It.isAnyString()))

.returns((e) => {

throw e;

});

appShell

.setup((a) => a.showInformationMessage(TypeMoq.It.isAny(), TypeMoq.It.isAny()))

.returns(() => Promise.resolve(''));

appShell

.setup((a) =>

a.showInformationMessage(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())

)

.returns((_a1: string, a2: string, _a3: string) => Promise.resolve(a2));

appShell

.setup((a) =>

a.showInformationMessage(

TypeMoq.It.isAny(),

TypeMoq.It.isAny(),

TypeMoq.It.isAny(),

TypeMoq.It.isAny()

)

)

.returns((_a1: string, a2: string, _a3: string, _a4: string) => Promise.resolve(a2));

appShell.setup((a) => a.showInputBox(TypeMoq.It.isAny())).returns(() => Promise.resolve(''));

appShell.setup((a) => a.setStatusBarMessage(TypeMoq.It.isAny())).returns(() => dummyDisposable);

ioc.serviceManager.rebindInstance<IApplicationShell>(IApplicationShell, appShell.object);

}

);

runTest('Remote Password', async () => {

const pythonService = await createPythonService();

if (pythonService && !useRawKernel) {

const configFile = path.join(

EXTENSION_ROOT_DIR,

'src',

'test',

'datascience',

'serverConfigFiles',

'remotePassword.py'

);

const uri = await startRemoteServer(pythonService, [

'-m',

'jupyter',

'notebook',

`--config=${configFile}`

]);

traceInfo('Waiting for notebook');

// We have a connection string here, so try to connect jupyterExecution to the notebook server

const notebook = await createNotebook(uri);

if (!notebook) {

assert.fail('Failed to connect to remote password server');

} else {

await verifySimple(notebook, `a=1${os.EOL}a`, 1);

}

}

});

runTest('Remote', async () => {

const pythonService = await createPythonService();

if (pythonService) {

const configFile = path.join(

EXTENSION_ROOT_DIR,

'src',

'test',

'datascience',

'serverConfigFiles',

'remoteToken.py'

);

const uri = await startRemoteServer(pythonService, [

'-m',

'jupyter',

'notebook',

`--config=${configFile}`

]);

// We have a connection string here, so try to connect jupyterExecution to the notebook server

const notebook = await createNotebook(uri);

if (!notebook) {

assert.fail('Failed to connect to remote server');

} else {

await verifySimple(notebook, `a=1${os.EOL}a`, 1);

}

}

});

runTest('Creation', async () => {

await createNotebook();

});

runTest('Failure', async (_this: Mocha.Context) => {

if (!useRawKernel) {

// Make a dummy class that will fail during launch

class FailedProcess extends JupyterExecutionFactory {

public isNotebookSupported = (): Promise<boolean> => {

return Promise.resolve(false);

};

}

ioc.serviceManager.rebind<IJupyterExecution>(IJupyterExecution, FailedProcess);

await createNotebook(undefined, undefined, true);

} else {

// This test is useless for raw kernel. You can't fail to launch a python process

_this.skip();

}

});

test('Not installed', async function () {

if (!useRawKernel) {

// Rewire our data we use to search for processes

@injectable()

class EmptyInterpreterService implements IInterpreterService {

public get hasInterpreters(): Promise<boolean> {

return Promise.resolve(true);

}

public onDidChangeInterpreter(

_listener: (e: void) => any,

_thisArgs?: any,

_disposables?: Disposable[]

): Disposable {

return { dispose: noop };

}

public onDidChangeInterpreterInformation(

_listener: (e: PythonInterpreter) => any,

_thisArgs?: any,

_disposables?: Disposable[]

): Disposable {

return { dispose: noop };

}

public getInterpreters(_resource?: Uri): Promise<PythonInterpreter[]> {

return Promise.resolve([]);

}

public autoSetInterpreter(): Promise<void> {

throw new Error('Method not implemented');

}

public getActiveInterpreter(_resource?: Uri): Promise<PythonInterpreter | undefined> {

return Promise.resolve(undefined);

}

public getInterpreterDetails(_pythonPath: string, _resoure?: Uri): Promise<PythonInterpreter> {

throw new Error('Method not implemented');

}

public refresh(_resource: Uri): Promise<void> {

throw new Error('Method not implemented');

}

public initialize(): void {

throw new Error('Method not implemented');

}

public getDisplayName(_interpreter: Partial<PythonInterpreter>): Promise<string> {

throw new Error('Method not implemented');

}

public shouldAutoSetInterpreter(): Promise<boolean> {

throw new Error('Method not implemented');

}

}

@injectable()

class EmptyPathService implements IKnownSearchPathsForInterpreters {

public getSearchPaths(): string[] {

return [];

}

}

ioc.serviceManager.rebind<IInterpreterService>(IInterpreterService, EmptyInterpreterService);

ioc.serviceManager.rebind<IKnownSearchPathsForInterpreters>(

IKnownSearchPathsForInterpreters,

EmptyPathService

);

await createNotebook(undefined, undefined, true);

} else {

// tslint:disable-next-line: no-invalid-this

this.skip();

}

});

runTest('Export/Import', async () => {

// Get a bunch of test cells (use our test cells from the react controls)

const testFolderPath = path.join(EXTENSION_ROOT_DIR, 'src', 'test', 'datascience');

const testState = generateTestState(testFolderPath);

const cells = testState.cellVMs.map((cellVM: ICellViewModel, _index: number) => {

return cellVM.cell;

});

// Translate this into a notebook

// Make sure we have a change dir happening

const settings = { ...ioc.getSettings().datascience };

settings.changeDirOnImportExport = true;

ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath, settings);

const exporter = ioc.serviceManager.get<INotebookExporter>(INotebookExporter);

const newFolderPath = path.join(

EXTENSION_ROOT_DIR,

'src',

'test',

'datascience',

'WorkspaceDir',

'WorkspaceSubDir',

'foo.ipynb'

);

const notebook = await exporter.translateToNotebook(cells, newFolderPath);

assert.ok(notebook, 'Translate to notebook is failing');

// Make sure we added in our chdir

if (notebook) {

const nbcells = notebook.cells;

if (nbcells) {

const firstCellText: string = nbcells[0].source as string;

assert.ok(firstCellText.includes('os.chdir'), `${firstCellText} does not include 'os.chdir`);

}

}

// Save to a temp file

const fileSystem = ioc.serviceManager.get<IFileSystem>(IFileSystem);

const importer = ioc.serviceManager.get<INotebookImporter>(INotebookImporter);

const temp = await fileSystem.createTemporaryFile('.ipynb');

try {

await fs.writeFile(temp.filePath, JSON.stringify(notebook), 'utf8');

// Try importing this. This should verify export works and that importing is possible

const results = await importer.importFromFile(temp.filePath);

// Make sure we have a single chdir in our results

const first = results.indexOf('os.chdir');

assert.ok(first >= 0, 'No os.chdir in import');

const second = results.indexOf('os.chdir', first + 1);

assert.equal(second, -1, 'More than one chdir in the import. It should be skipped');

// Make sure we have a cell in our results

assert.ok(/#\s*%%/.test(results), 'No cells in returned import');

} finally {

importer.dispose();

temp.dispose();

}

});

// tslint:disable-next-line:no-invalid-template-strings

runTest('Verify ${fileDirname} working directory', async () => {

// Verify that the default ${fileDirname} setting sets the working directory to the file path

changeMockWorkingDirectory(`'${srcDirectory()}'`);

const notebook = await createNotebook();

await verifySimple(notebook, 'import os\nos.getcwd()', srcDirectory(), true);

await verifySimple(notebook, 'import sys\nsys.path[0]', srcDirectory(), true);

});

runTest('Change Interpreter', async () => {

const isRollingBuild = process.env ? process.env.VSCODE_PYTHON_ROLLING !== undefined : false;

// Real Jupyter doesn't help this test at all and is tricky to set up for it, so just skip it

if (!isRollingBuild) {

const server = await createNotebook();

// Create again, we should get the same server from the cache

const server2 = await createNotebook();

// tslint:disable-next-line: triple-equals

assert.ok(server == server2, 'With no settings changed we should return the cached server');

// Create a new mock interpreter with a different path

const newPython: PythonInterpreter = {

path: '/foo/bar/baz/python.exe',

version: new SemVer('3.6.6-final'),

sysVersion: '1.0.0.0',

sysPrefix: 'Python',

type: InterpreterType.Unknown,

architecture: Architecture.x64

};

// Add interpreter into mock jupyter service and set it as active

ioc.addInterpreter(newPython, SupportedCommands.all);

// Create a new notebook, we should still be the same as interpreter is just saved for notebook creation

const server3 = await createNotebook();

// tslint:disable-next-line: triple-equals

assert.ok(server == server3, 'With interpreter changed we should not return a new server');

} else {

console.log(`Skipping Change Interpreter test in non-mocked Jupyter case`);

}

});

runTest('Restart kernel', async () => {

addMockData(`a=1${os.EOL}a`, 1);

addMockData(`a+=1${os.EOL}a`, 2);

addMockData(`a+=4${os.EOL}a`, 6);

addMockData('a', `name 'a' is not defined`, 'error');

const server = await createNotebook();

// Setup some state and verify output is correct

await verifySimple(server, `a=1${os.EOL}a`, 1);

await verifySimple(server, `a+=1${os.EOL}a`, 2);

await verifySimple(server, `a+=4${os.EOL}a`, 6);

console.log('Waiting for idle');

// In unit tests we have to wait for status idle before restarting. Unit tests

// seem to be timing out if the restart throws any exceptions (even if they're caught)

await server!.waitForIdle(10000);

console.log('Restarting kernel');

try {

await server!.restartKernel(10000);

console.log('Waiting for idle');

await server!.waitForIdle(10000);

console.log('Verifying restart');

await verifyError(server, 'a', `name 'a' is not defined`);

} catch (exc) {

assert.ok(

exc instanceof JupyterKernelPromiseFailedError,

`Restarting did not timeout correctly for ${exc}`

);

}

});

class TaggedCancellationTokenSource extends CancellationTokenSource {

public tag: string;

constructor(tag: string) {

super();

this.tag = tag;

}

}

async function testCancelableCall<T>(

method: (t: CancellationToken) => Promise<T>,

messageFormat: string,

timeout: number

): Promise<boolean> {

const tokenSource = new TaggedCancellationTokenSource(messageFormat.format(timeout.toString()));

const disp = setTimeout(

(_s) => {

tokenSource.cancel();

},

timeout,

tokenSource.tag

);

try {

// tslint:disable-next-line:no-string-literal

(tokenSource.token as any)['tag'] = messageFormat.format(timeout.toString());

await method(tokenSource.token);

} catch (exc) {

// This should happen. This means it was canceled.

assert.ok(exc instanceof CancellationError, `Non cancellation error found : ${exc.stack}`);

} finally {

clearTimeout(disp);

tokenSource.dispose();

}

return true;

}

async function testCancelableMethod<T>(

method: (t: CancellationToken) => Promise<T>,

messageFormat: string,

short?: boolean

): Promise<boolean> {

const timeouts = short ? [10, 20, 30, 100] : [300, 400, 500, 1000];

// tslint:disable-next-line:prefer-for-of

for (let i = 0; i < timeouts.length; i += 1) {

await testCancelableCall(method, messageFormat, timeouts[i]);

}

return true;

}

runTest('Cancel execution', async (_this: Mocha.Context) => {

if (useRawKernel) {

// Not cancellable at the moment. Just starts a process

_this.skip();

return;

}

if (ioc.mockJupyter) {

ioc.mockJupyter.setProcessDelay(2000);

addMockData(`a=1${os.EOL}a`, 1);

}

const jupyterExecution = ioc.get<IJupyterExecution>(IJupyterExecution);

// Try different timeouts, canceling after the timeout on each

assert.ok(

await testCancelableMethod(

(t: CancellationToken) => jupyterExecution.connectToNotebookServer(undefined, t),

'Cancel did not cancel start after {0}ms'

)

);

if (ioc.mockJupyter) {

ioc.mockJupyter.setProcessDelay(undefined);

}

// Make sure doing normal start still works

const nonCancelSource = new CancellationTokenSource();

const server = await jupyterExecution.connectToNotebookServer(undefined, nonCancelSource.token);

const notebook = server

? await server.createNotebook(baseUri, Uri.parse(Identifiers.InteractiveWindowIdentity))

: undefined;

assert.ok(notebook, 'Server not found with a cancel token that does not cancel');

// Make sure can run some code too

await verifySimple(notebook, `a=1${os.EOL}a`, 1);

if (ioc.mockJupyter) {

ioc.mockJupyter.setProcessDelay(200);

}

// Force a settings changed so that all of the cached data is cleared

ioc.forceSettingsChanged(undefined, '/usr/bin/test3/python');

assert.ok(

await testCancelableMethod(

(t: CancellationToken) => jupyterExecution.getUsableJupyterPython(t),

'Cancel did not cancel getusable after {0}ms',

true

)

);

assert.ok(

await testCancelableMethod(

(t: CancellationToken) => jupyterExecution.isNotebookSupported(t),

'Cancel did not cancel isNotebook after {0}ms',

true

)

);

assert.ok(

await testCancelableMethod(

(t: CancellationToken) => jupyterExecution.isImportSupported(t),

'Cancel did not cancel isImport after {0}ms',

true

)

);

});

async function interruptExecute(

notebook: INotebook | undefined,

code: string,

interruptMs: number,

sleepMs: number

): Promise<InterruptResult> {

let interrupted = false;

let finishedBefore = false;

const finishedPromise = createDeferred();

let error;

const observable = notebook!.executeObservable(code, Uri.file('foo.py').fsPath, 0, uuid(), false);

observable.subscribe(

(c) => {

if (c.length > 0 && c[0].state === CellState.error) {

finishedBefore = !interrupted;

finishedPromise.resolve();

}

if (c.length > 0 && c[0].state === CellState.finished) {

finishedBefore = !interrupted;

finishedPromise.resolve();

}

},

(err) => {

error = err;

finishedPromise.resolve();

},

() => finishedPromise.resolve()

);

// Then interrupt

interrupted = true;

const result = await notebook!.interruptKernel(interruptMs);

// Then we should get our finish unless there was a restart

await waitForPromise(finishedPromise.promise, sleepMs);

assert.equal(finishedBefore, false, 'Finished before the interruption');

assert.equal(error, undefined, 'Error thrown during interrupt');

assert.ok(

finishedPromise.completed ||

result === InterruptResult.TimedOut ||

result === InterruptResult.Success,

`Interrupt restarted ${result} for: ${code}`

);

return result;

}

runTest('Interrupt kernel', async (_this: Mocha.Context) => {

// Interrupt doesn't work yet for the raw kernel.

if (useRawKernel) {

_this.skip();

return;

}

const returnable = `import signal

import _thread

import time

keep_going = True

def handler(signum, frame):

global keep_going

print('signal')

keep_going = False

signal.signal(signal.SIGINT, handler)

while keep_going:

print(".")

time.sleep(.1)`;

const fourSecondSleep = `import time${os.EOL}time.sleep(4)${os.EOL}print("foo")`;

const kill = `import signal

import time

import os

keep_going = True

def handler(signum, frame):