Skip to content

Commit

Permalink
refine test session shutdown & restart behavior; fix related tests; s…
Browse files Browse the repository at this point in the history 
…tart queuing tests
  • Loading branch information
@seeM
seeM committed Nov 19, 2024
1 parent e3eedc8 commit 9e1e2dd
Show file tree
Hide file tree
Showing 2 changed files with 147 additions and 89 deletions.
205 changes: 117 additions & 88 deletions src/vs/workbench/services/runtimeSession/test/common/runtimeSession.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

import { strict as assert } from 'assert';
import * as sinon from 'sinon';
import { CancellationError } from 'vs/base/common/errors';
import { Event } from 'vs/base/common/event';
import { URI } from 'vs/base/common/uri';
import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils';
Expand All @@ -15,7 +14,7 @@ import { TestInstantiationService } from 'vs/platform/instantiation/test/common/
import { IWorkspaceTrustManagementService } from 'vs/platform/workspace/common/workspaceTrust';
import { formatLanguageRuntimeMetadata, ILanguageRuntimeMetadata, ILanguageRuntimeService, LanguageRuntimeSessionMode, RuntimeExitReason, RuntimeState } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService';
import { ILanguageRuntimeSession, IRuntimeSessionMetadata, IRuntimeSessionService, IRuntimeSessionWillStartEvent } from 'vs/workbench/services/runtimeSession/common/runtimeSessionService';
import { TestLanguageRuntimeSession } from 'vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession';
import { TestLanguageRuntimeSession, waitForRuntimeState } from 'vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession';
import { createRuntimeServices, createTestLanguageRuntimeMetadata, startTestLanguageRuntimeSession } from 'vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService';
import { TestRuntimeSessionManager } from 'vs/workbench/test/common/positronWorkbenchTestServices';
import { TestWorkspaceTrustManagementService } from 'vs/workbench/test/common/workbenchTestServices';
Expand Down Expand Up @@ -156,16 +155,29 @@ suite('Positron - RuntimeSessionService', () => {
}
}

function assertSessionIsStarting(session: ILanguageRuntimeSession) {
assert.equal(session.getRuntimeState(), RuntimeState.Starting);

function assertSingleSession(session: ILanguageRuntimeSession) {
if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) {
assertServiceState({ hasStartingOrRunningConsole: true, consoleSession: session }, session.runtimeMetadata);
} else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) {
assertServiceState({ notebookSession: session, notebookSessionForNotebookUri: session }, session.runtimeMetadata);
}
}

function assertSessionIsStarting(session: ILanguageRuntimeSession) {
assert.equal(session.getRuntimeState(), RuntimeState.Starting);
assertSingleSession(session);
}

function assertSessionIsRestarting(session: ILanguageRuntimeSession) {
assert.equal(session.getRuntimeState(), RuntimeState.Restarting);
assertSingleSession(session);
}

function assertSessionIsStarted(session: ILanguageRuntimeSession) {
assert.equal(session.getRuntimeState(), RuntimeState.Ready);
assertSingleSession(session);
}

async function restoreSession(
sessionMetadata: IRuntimeSessionMetadata, runtimeMetadata = runtime,
) {
Expand Down Expand Up @@ -568,6 +580,7 @@ suite('Positron - RuntimeSessionService', () => {

test('select console while another runtime is running for the language', async () => {
const session1 = await startConsole(anotherRuntime);
await waitForRuntimeState(session1, RuntimeState.Ready);
const session2 = await selectRuntime();

assert.equal(session1.getRuntimeState(), RuntimeState.Exited);
Expand All @@ -580,6 +593,14 @@ suite('Positron - RuntimeSessionService', () => {
});
});

test('select console throws if session is still starting', async () => {
await startConsole(anotherRuntime);
await assert.rejects(
selectRuntime(),
new Error('Cannot shut down kernel; it is not (yet) running. (state = starting)'),
);
});

test('select console to the same runtime sets the foreground session', async () => {
const session1 = await startConsole();

Expand Down Expand Up @@ -620,14 +641,19 @@ suite('Positron - RuntimeSessionService', () => {

await restartSession(session.sessionId);

assertSessionIsStarting(session);
assert.equal(session.getRuntimeState(), RuntimeState.Restarting);
assertSessionIsRestarting(session);

await waitForRuntimeState(session, RuntimeState.Ready);
assertSessionIsStarted(session);
});
}

for (const state of [RuntimeState.Uninitialized, RuntimeState.Exited]) {
test(`restart ${mode} in '${state}' state`, async () => {
// Get a session to the exited state.
const session = await start();
await waitForRuntimeState(session, RuntimeState.Ready);
await session.shutdown(RuntimeExitReason.Shutdown);
await waitForRuntimeState(session, RuntimeState.Exited);

Expand Down Expand Up @@ -688,7 +714,7 @@ suite('Positron - RuntimeSessionService', () => {

await restartSession(session.sessionId);

assertSessionIsStarting(session);
assertSessionIsRestarting(session);

sinon.assert.notCalled(target);
});
Expand All @@ -705,7 +731,7 @@ suite('Positron - RuntimeSessionService', () => {
restartSession(session.sessionId),
]);

assertSessionIsStarting(session);
assertSessionIsRestarting(session);

sinon.assert.calledOnce(target);
});
Expand All @@ -722,91 +748,94 @@ suite('Positron - RuntimeSessionService', () => {
await waitForRuntimeState(session, RuntimeState.Ready);
await restartSession(session.sessionId);

assertSessionIsStarting(session);
assertSessionIsRestarting(session);

sinon.assert.calledThrice(target);
});
}

// test(`${name} notebook while shutting down`, async () => {
// const session1 = await startNotebook();

// const [, session2,] = await Promise.all([
// runtimeSessionService.shutdownNotebookSession(notebookUri, RuntimeExitReason.Shutdown),
// startNotebook(),
// ]);

// assert.equal(session1.getRuntimeState(), RuntimeState.Exited);
// assert.equal(session2.getRuntimeState(), RuntimeState.Starting);
// assertServiceState({
// notebookSession: session2,
// notebookSessionForNotebookUri: session2,
// activeSessions: [session1, session2],
// });
// });

// test(`${name} notebook while restarting and in "exited" state`, async () => {
// const session = await startNotebook();
// await waitForRuntimeState(session, RuntimeState.Ready);

// const target = sinon.spy(session, 'restart');

// const startPromise = new Promise<TestLanguageRuntimeSession>(resolve => {
// const disposable = session.onDidChangeRuntimeState(state => {
// if (state === RuntimeState.Exited) {
// disposable.dispose();
// resolve(startNotebook());
// }
// });
// });

// const [, session2,] = await Promise.all([
// restartSession(session.sessionId),
// startPromise,
// // startNotebook(),
// ]);

// assert.equal(session, session2);

// assert.equal(session.getRuntimeState(), RuntimeState.Starting);
// assertServiceState({
// notebookSession: session,
// notebookSessionForNotebookUri: session,
// });

// sinon.assert.calledOnce(target);
// });

// test('restart notebook while shutting down', async () => {
// const session = await startNotebook();

// await Promise.all([
// runtimeSessionService.shutdownNotebookSession(notebookUri, RuntimeExitReason.Shutdown),
// restartSession(session.sessionId),
// ]);

// assert.equal(session.getRuntimeState(), RuntimeState.Starting);
// assertNotebookSessionIsStarted(session);
// });
});
suite('queueing', () => {
// Consider the following sequence of actions:
// start -> restart == should do both -- not possible since restart requires session ID
// start -> select == select should noop
// restart -> start == start should noop
// restart -> select == select should noop
// select -> start == start should noop
// select -> restart == restart should noop -- not possible since restart requires session ID
test(`start console -> select console`, async () => {
const [session,] = await Promise.all([
startConsole(),
selectRuntime(),
]);

async function waitForRuntimeState(
session: ILanguageRuntimeSession,
state: RuntimeState,
timeout = 10_000,
) {
return new Promise<void>((resolve, reject) => {
const timer = setTimeout(() => {
disposable.dispose();
reject(new CancellationError());
}, timeout);

const disposable = session.onDidChangeRuntimeState(newState => {
if (newState === state) {
clearTimeout(timer);
disposable.dispose();
resolve();
}
assertSessionIsStarting(session);
});

// TODO: I think this behavior is undefined at the moment... We don't have any safety checks
// in start based on restart behavior.
test(`restart console -> start console`, async () => {
const session = await startConsole();
await waitForRuntimeState(session, RuntimeState.Ready);

const didChangeRuntimeState = sinon.spy((_state: RuntimeState) => { });
disposables.add(session.onDidChangeRuntimeState(didChangeRuntimeState));

const promise2 = new Promise<ILanguageRuntimeSession>(resolve => {
disposables.add(session.onDidChangeRuntimeState(state => {
if (state === RuntimeState.Exited) {
const result = startConsole();
resolve(result);
}
}));
});

const promise = restartSession(session.sessionId);

await Promise.all([promise, promise2]);

await waitForRuntimeState(session, RuntimeState.Ready);

console.log(didChangeRuntimeState.getCalls().map(call => call.args));
sinon.assert.callCount(didChangeRuntimeState, 4);
sinon.assert.callOrder(
didChangeRuntimeState.withArgs(RuntimeState.Restarting),
didChangeRuntimeState.withArgs(RuntimeState.Exited),
didChangeRuntimeState.withArgs(RuntimeState.Starting),
didChangeRuntimeState.withArgs(RuntimeState.Ready),
);
sinon.assert.calledOnce(didChangeRuntimeState.withArgs(RuntimeState.Exited));

// await Promise.all([
// restartSession(session.sessionId),
// // TODO: This returns the existing sessionId because it's still considered 'running'.
// // Shouldn't the session going to 'exited' remove it from console sessions?
// startConsole(),
// ]);

// TODO: Is it possible to start a session immediately when the session goes to 'exited'?

assertSessionIsRestarting(session);
});

// test(`restart console -> select console`, async () => {
// const session1 = await startConsole();
// await waitForRuntimeState(session1, RuntimeState.Ready);

// const target = sinon.spy((_state: RuntimeState) => { });
// disposables.add(session1.onDidChangeRuntimeState(target));

// const [, session2] = await Promise.all([
// restartSession(session1.sessionId),
// selectRuntime(),
// ]);

// assert.equal(session1.sessionId, session2.sessionId);

// assertSessionIsStarting(session1);

// sinon.assert.called(target);
// sinon.assert.calledOnce(target.withArgs(RuntimeState.Starting));
// console.log(target.getCalls().map(call => call.args));
// });
});
}
});
31 changes: 30 additions & 1 deletion src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import { ILanguageRuntimeSession, IRuntimeClientInstance, IRuntimeSessionMetadat
import { ILanguageRuntimeClientCreatedEvent, ILanguageRuntimeExit, ILanguageRuntimeInfo, ILanguageRuntimeMessage, ILanguageRuntimeMessageClearOutput, ILanguageRuntimeMessageError, ILanguageRuntimeMessageInput, ILanguageRuntimeMessageIPyWidget, ILanguageRuntimeMessageOutput, ILanguageRuntimeMessagePrompt, ILanguageRuntimeMessageResult, ILanguageRuntimeMessageState, ILanguageRuntimeMessageStream, ILanguageRuntimeMetadata, ILanguageRuntimeStartupFailure, LanguageRuntimeMessageType, RuntimeCodeExecutionMode, RuntimeCodeFragmentStatus, RuntimeErrorBehavior, RuntimeExitReason, RuntimeOnlineState, RuntimeOutputKind, RuntimeState } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService';
import { IRuntimeClientEvent } from 'vs/workbench/services/languageRuntime/common/languageRuntimeUiClient';
import { TestRuntimeClientInstance } from 'vs/workbench/services/languageRuntime/test/common/testRuntimeClientInstance';
import { CancellationError } from 'vs/base/common/errors';

export class TestLanguageRuntimeSession extends Disposable implements ILanguageRuntimeSession {
private readonly _onDidChangeRuntimeState = this._register(new Emitter<RuntimeState>());
Expand Down Expand Up @@ -162,10 +163,17 @@ export class TestLanguageRuntimeSession extends Disposable implements ILanguageR

async restart(): Promise<void> {
await this.shutdown(RuntimeExitReason.Restart);
await this.start();
waitForRuntimeState(this, RuntimeState.Exited).then(() => this.start());
}

async shutdown(exitReason: RuntimeExitReason): Promise<void> {
if (this._currentState !== RuntimeState.Idle &&
this._currentState !== RuntimeState.Busy &&
this._currentState !== RuntimeState.Ready) {
throw new Error('Cannot shut down kernel; it is not (yet) running.' +
` (state = ${this._currentState})`);
}

if (exitReason === RuntimeExitReason.Restart) {
this._onDidChangeRuntimeState.fire(RuntimeState.Restarting);
} else {
Expand Down Expand Up @@ -362,3 +370,24 @@ export class TestLanguageRuntimeSession extends Disposable implements ILanguageR
});
}
}

export async function waitForRuntimeState(
session: ILanguageRuntimeSession,
state: RuntimeState,
timeout = 10_000,
) {
return new Promise<void>((resolve, reject) => {
const timer = setTimeout(() => {
disposable.dispose();
reject(new CancellationError());
}, timeout);

const disposable = session.onDidChangeRuntimeState(newState => {
if (newState === state) {
clearTimeout(timer);
disposable.dispose();
resolve();
}
});
});
}

0 comments on commit 9e1e2dd

Please sign in to comment.