Skip to content

Commit

Permalink
Add necessary tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@DonJayamanne
DonJayamanne committed Dec 8, 2021
1 parent ebaaea7 commit 2fa9f8b
Show file tree
Hide file tree
Showing 3 changed files with 190 additions and 16 deletions.
21 changes: 6 additions & 15 deletions src/client/datascience/notebook/notebookControllerManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -290,23 +290,14 @@ export class NotebookControllerManager implements INotebookControllerManager, IE
let defaultPythonKernel: VSCodeNotebookController | undefined;
let defaultPythonLanguageKernel: VSCodeNotebookController | undefined;
controllers.forEach((item) => {
if (
(item.connection.kind === 'startUsingLocalKernelSpec' ||
item.connection.kind === 'startUsingRemoteKernelSpec') &&
item.connection.kernelSpec.name === 'python'
) {
if (item.connection.kind !== 'startUsingRemoteKernelSpec') {
return;
}
if (item.connection.kernelSpec.name === 'python') {
defaultPythonKernel = item;
} else if (
(item.connection.kind === 'startUsingLocalKernelSpec' ||
item.connection.kind === 'startUsingRemoteKernelSpec') &&
item.connection.kernelSpec.name === 'python3'
) {
} else if (item.connection.kernelSpec.name === 'python3') {
defaultPython3Kernel = item;
} else if (
(item.connection.kind === 'startUsingLocalKernelSpec' ||
item.connection.kind === 'startUsingRemoteKernelSpec') &&
item.connection.kernelSpec.language === PYTHON_LANGUAGE
) {
} else if (item.connection.kernelSpec.language === PYTHON_LANGUAGE) {
defaultPythonLanguageKernel = item;
}
});
Expand Down
35 changes: 35 additions & 0 deletions src/test/datascience/jupyterServer.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,20 @@ export class JupyterServer implements IAsyncDisposable {
private static _instance: JupyterServer;
private _disposables: IDisposable[] = [];
private _jupyterServerWithToken?: Promise<Uri>;
private _secondJupyterServerWithToken?: Promise<Uri>;
private availablePort?: number;
private availableSecondPort?: number;
public async dispose() {
this._jupyterServerWithToken = undefined;
this._secondJupyterServerWithToken = undefined;
disposeAllDisposables(this._disposables);
traceInfo('Shutting Jupyter server used for remote tests');
if (this.availablePort) {
await tcpPortUsed.waitUntilFree(this.availablePort, 200, 5_000);
}
if (this.availableSecondPort) {
await tcpPortUsed.waitUntilFree(this.availableSecondPort, 200, 5_000);
}
}
public async startJupyterWithToken(token = '7d25707a86975be50ee9757c929fef9012d27cf43153d1c1'): Promise<Uri> {
if (!this._jupyterServerWithToken) {
Expand All @@ -55,6 +61,27 @@ export class JupyterServer implements IAsyncDisposable {
}
return this._jupyterServerWithToken;
}
public async startSecondJupyterWithToken(token = 'fbd00a866c54f5d9f64df9ba820860de56f32379407d03e8'): Promise<Uri> {
if (!this._secondJupyterServerWithToken) {
this._secondJupyterServerWithToken = new Promise<Uri>(async (resolve, reject) => {
const port = await this.getSecondFreePort();
// Possible previous instance of jupyter has not completely shutdown.
// Wait for it to shutdown fully so that we can re-use the same port.
await tcpPortUsed.waitUntilFree(port, 200, 10_000);
try {
await this.startJupyterServer({
port,
token
});
await sleep(5_000); // Wait for some time for Jupyter to warm up & be ready to accept connections.
resolve(Uri.parse(`http://localhost:${port}/?token=${token}`));
} catch (ex) {
reject(ex);
}
});
}
return this._secondJupyterServerWithToken;
}
private async getFreePort() {
// Always use the same port (when using different ports, our code doesn't work as we need to re-load VSC).
// The remote uri is cached in a few places (known issue).
Expand All @@ -63,6 +90,14 @@ export class JupyterServer implements IAsyncDisposable {
}
return this.availablePort!;
}
private async getSecondFreePort() {
// Always use the same port (when using different ports, our code doesn't work as we need to re-load VSC).
// The remote uri is cached in a few places (known issue).
if (!this.availableSecondPort) {
this.availableSecondPort = await getFreePort({ host: 'localhost' }).then((p) => p);
}
return this.availableSecondPort!;
}

private startJupyterServer({ token, port }: { token: string; port: number }): Promise<void> {
return new Promise<void>(async (resolve, reject) => {
Expand Down
150 changes: 149 additions & 1 deletion src/test/datascience/notebook/remoteNotebookEditor.vscode.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,19 @@ import {
runCell,
deleteAllCellsAndWait,
insertCodeCell,
waitForTextOutput
waitForTextOutput,
defaultNotebookTestTimeout,
createEmptyPythonNotebook
} from './helper';
import { openNotebook } from '../helpers';
import { PYTHON_LANGUAGE } from '../../../client/common/constants';
import { PreferredRemoteKernelIdProvider } from '../../../client/datascience/notebookStorage/preferredRemoteKernelIdProvider';
import { Settings } from '../../../client/datascience/constants';
import { INotebookControllerManager } from '../../../client/datascience/notebook/types';
import { JupyterServerSelector } from '../../../client/datascience/jupyter/serverSelector';
import { RemoteKernelSpecConnectionMetadata } from '../../../client/datascience/jupyter/kernels/types';
import { JupyterServer } from '../jupyterServer';
import { JVSC_EXTENSION_ID_FOR_TESTS } from '../../constants';

/* eslint-disable @typescript-eslint/no-explicit-any, no-invalid-this */
suite('DataScience - VSCode Notebook - (Remote) (Execution) (slow)', function () {
Expand All @@ -50,6 +57,9 @@ suite('DataScience - VSCode Notebook - (Remote) (Execution) (slow)', function ()
let ipynbFile: Uri;
let globalMemento: Memento;
let encryptedStorage: IEncryptedStorage;
let controllerManager: INotebookControllerManager;
let jupyterServerSelector: JupyterServerSelector;

suiteSetup(async function () {
if (!IS_REMOTE_NATIVE_TEST) {
return this.skip();
Expand All @@ -61,9 +71,14 @@ suite('DataScience - VSCode Notebook - (Remote) (Execution) (slow)', function ()
}
await startJupyterServer();
sinon.restore();
jupyterServerSelector = api.serviceContainer.get<JupyterServerSelector>(JupyterServerSelector);
vscodeNotebook = api.serviceContainer.get<IVSCodeNotebook>(IVSCodeNotebook);
encryptedStorage = api.serviceContainer.get<IEncryptedStorage>(IEncryptedStorage);
globalMemento = api.serviceContainer.get<Memento>(IMemento, GLOBAL_MEMENTO);
controllerManager = api.serviceContainer.get<INotebookControllerManager>(
INotebookControllerManager,
INotebookControllerManager
);
remoteKernelIdProvider = api.serviceContainer.get<PreferredRemoteKernelIdProvider>(
PreferredRemoteKernelIdProvider
);
Expand Down Expand Up @@ -158,4 +173,137 @@ suite('DataScience - VSCode Notebook - (Remote) (Execution) (slow)', function ()
waitForTextOutput(cell2, 'Hello World', 0, false)
]);
});
test('Local and Remote kernels are listed', async function () {
await controllerManager.loadNotebookControllers();
const controllers = await controllerManager.kernelConnections;
assert.ok(
controllers.some((item) => item.kind === 'startUsingRemoteKernelSpec'),
'Should have at least one remote kernelspec'
);
assert.ok(
controllers.some(
(item) => item.kind === 'startUsingLocalKernelSpec' || item.kind === 'startUsingPythonInterpreter'
),
'Should have at least one local kernel'
);
});
test('Remote kernels are removed when switching to local', async function () {
await controllerManager.loadNotebookControllers();
assert.ok(async () => {
const controllers = await controllerManager.kernelConnections;
return controllers.filter((item) => item.kind === 'startUsingRemoteKernelSpec').length === 0;
}, 'Should have at least one remote kernelspec');

// After resetting connection to local only, then remove all remote connections.
await jupyterServerSelector.setJupyterURIToLocal();

await waitForCondition(
async () => {
const controllers = await controllerManager.kernelConnections;
return controllers.filter((item) => item.kind === 'startUsingRemoteKernelSpec').length === 0;
},
defaultNotebookTestTimeout,
'Should not have any remote controllers'
);
});

test('Old Remote kernels are removed when switching to new Remote Server', async function () {
await controllerManager.loadNotebookControllers();
const baseUrls = new Set<string>();
const controllers = await controllerManager.kernelConnections;
const remoteKernelSpecs = controllers
.filter((item) => item.kind === 'startUsingRemoteKernelSpec')
.map((item) => item as RemoteKernelSpecConnectionMetadata);
remoteKernelSpecs.forEach((item) => baseUrls.add(item.baseUrl));
assert.isOk(remoteKernelSpecs.length > 0, 'Should have at least one remote kernelspec');

// Start another jupyter server with a new port.
const uri = await JupyterServer.instance.startJupyterWithToken();
const uriString = decodeURIComponent(uri.toString());
traceInfo(`Jupyter started and listening at ${uriString}`);
await jupyterServerSelector.setJupyterURIToRemote(uriString);

// Wait til we get new controllers with a different base url.
await waitForCondition(
async () => {
const controllers = await controllerManager.kernelConnections;
return controllers.some(
(item) => item.kind === 'startUsingRemoteKernelSpec' && !baseUrls.has(item.baseUrl)
);
},
defaultNotebookTestTimeout,
'Should have at least one remote kernelspec with different baseUrls'
);
});

test('Local Kernel state is not lost when connecting to remote', async function () {
await controllerManager.loadNotebookControllers();

// After resetting connection to local only, verify all remote connections are no longer available.
await jupyterServerSelector.setJupyterURIToLocal();
await waitForCondition(
async () => {
const controllers = await controllerManager.kernelConnections;
return controllers.filter((item) => item.kind === 'startUsingRemoteKernelSpec').length === 0;
},
defaultNotebookTestTimeout,
'Should not have any remote controllers'
);

await createEmptyPythonNotebook(disposables);
await insertCodeCell('a = "123412341234"', { index: 0 });
await insertCodeCell('print(a)', { index: 1 });
const cell1 = vscodeNotebook.activeNotebookEditor?.document.cellAt(0)!;
const cell2 = vscodeNotebook.activeNotebookEditor?.document.cellAt(0)!;
await runCell(cell1);

// Now that we don't have any remote kernels, connect to a remote jupyter server.
await startJupyterServer();

// Verify we have a remote kernel spec.
await waitForCondition(
async () => {
const controllers = await controllerManager.kernelConnections;
return controllers.some((item) => item.kind === 'startUsingRemoteKernelSpec');
},
defaultNotebookTestTimeout,
'Should have at least one remote controller'
);

// Run the second cell and verify we still have the same kernel state.
await Promise.all([runCell(cell2), waitForTextOutput(cell2, '123412341234')]);
});

test('Can run against a remtoe kernelspec', async function () {
await controllerManager.loadNotebookControllers();
const controllers = await controllerManager.kernelConnections;

// Verify we have a remote kernel spec.
assert.ok(
controllers.some((item) => item.kind === 'startUsingRemoteKernelSpec'),
'Should have at least one remote controller'
);

await createEmptyPythonNotebook(disposables);

// Find the default remote Python kernel (we know that will have ipykernel, as we've set up CI as such).
const defaultPythonKernel = await controllerManager.getActiveInterpreterOrDefaultController(
'jupyter-notebook',
undefined
);
assert.ok(defaultPythonKernel, 'No default remote kernel');
assert.strictEqual(
defaultPythonKernel?.connection.kind,
'startUsingRemoteKernelSpec',
'Not a remote kernelspec'
);
await this.commandManager.executeCommand('notebook.selectKernel', {
id: defaultPythonKernel!.controller.id,
extension: JVSC_EXTENSION_ID_FOR_TESTS
});

await insertCodeCell('print("123412341234")', { index: 0 });
const cell = vscodeNotebook.activeNotebookEditor?.document.cellAt(0)!;
await Promise.all([runCell(cell), waitForTextOutput(cell, '123412341234')]);
});
});

0 comments on commit 2fa9f8b

Please sign in to comment.