testing: allow work after user interrupts test (#168212)

Gives a 10 seconds timeout for test extensions to tear down and write
out remaining messages, instead of immediately finalizing the run when
the user clicks the stop button. Users can also click the test button
again to forcefully end the run.

src/vs/workbench/api/common/extHostTesting.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { mapFind } from 'vs/base/common/arrays';
+import { RunOnceScheduler } from 'vs/base/common/async';
 import { VSBuffer } from 'vs/base/common/buffer';
 import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
 import { Emitter, Event } from 'vs/base/common/event';
@@ -300,7 +301,7 @@ export class ExtHostTesting implements ExtHostTestingShape {
 		} catch (e) {
 			return { error: String(e) };
 		} finally {
-			if (tracker.isRunning && !token.isCancellationRequested) {
+			if (tracker.hasRunningTasks && !token.isCancellationRequested) {
 				await Event.toPromise(tracker.onEnd);
 			}
 
@@ -320,12 +321,24 @@ export class ExtHostTesting implements ExtHostTestingShape {
 	}
 }
 
+// Deadline after being requested by a user that a test run is forcibly cancelled.
+const RUN_CANCEL_DEADLINE = 10_000;
+
+const enum TestRunTrackerState {
+	// Default state
+	Running,
+	// Cancellation is requested, but the run is still going.
+	Cancelling,
+	// All tasks have ended
+	Ended,
+}
+
 class TestRunTracker extends Disposable {
+	private state = TestRunTrackerState.Running;
 	private readonly tasks = new Map</* task ID */string, { run: vscode.TestRun; coverage: TestRunCoverageBearer }>();
 	private readonly sharedTestIds = new Set<string>();
 	private readonly cts: CancellationTokenSource;
 	private readonly endEmitter = this._register(new Emitter<void>());
-	private disposed = false;
 
 	/**
 	 * Fires when a test ends, and no more tests are left running.
@@ -335,7 +348,7 @@ class TestRunTracker extends Disposable {
 	/**
 	 * Gets whether there are any tests running.
 	 */
-	public get isRunning() {
+	public get hasRunningTasks() {
 		return this.tasks.size > 0;
 	}
 
@@ -349,18 +362,28 @@ class TestRunTracker extends Disposable {
 	constructor(private readonly dto: TestRunDto, private readonly proxy: MainThreadTestingShape, parentToken?: CancellationToken) {
 		super();
 		this.cts = this._register(new CancellationTokenSource(parentToken));
-		this._register(this.cts.token.onCancellationRequested(() => {
-			for (const { run } of this.tasks.values()) {
-				run.end();
-			}
-		}));
+
+		const forciblyEnd = this._register(new RunOnceScheduler(() => this.forciblyEndTasks(), RUN_CANCEL_DEADLINE));
+		this._register(this.cts.token.onCancellationRequested(() => forciblyEnd.schedule()));
 	}
 
+	/** Requests cancellation of the run. On the second call, forces cancellation. */
+	public cancel() {
+		if (this.state === TestRunTrackerState.Running) {
+			this.cts.cancel();
+			this.state = TestRunTrackerState.Cancelling;
+		} else if (this.state === TestRunTrackerState.Cancelling) {
+			this.forciblyEndTasks();
+		}
+	}
+
+	/** Gets coverage for a task ID. */
 	public getCoverage(taskId: string) {
 		return this.tasks.get(taskId)?.coverage;
 	}
 
-	public createRun(name: string | undefined) {
+	/** Creates the public test run interface to give to extensions. */
+	public createRun(name: string | undefined): vscode.TestRun {
 		const runId = this.dto.id;
 		const ctrlId = this.dto.controllerId;
 		const taskId = generateUuid();
@@ -458,8 +481,8 @@ class TestRunTracker extends Disposable {
 				ended = true;
 				this.proxy.$finishedTestRunTask(runId, taskId);
 				this.tasks.delete(taskId);
-				if (!this.isRunning) {
-					this.dispose();
+				if (!this.tasks.size) {
+					this.markEnded();
 				}
 			}
 		};
@@ -470,15 +493,18 @@ class TestRunTracker extends Disposable {
 		return run;
 	}
 
-	public override dispose() {
-		if (!this.disposed) {
-			this.disposed = true;
-			this.endEmitter.fire();
-			this.cts.cancel();
-			super.dispose();
+	private forciblyEndTasks() {
+		for (const { run } of this.tasks.values()) {
+			run.end();
 		}
 	}
 
+	private markEnded() {
+		if (this.state !== TestRunTrackerState.Ended) {
+			this.state = TestRunTrackerState.Ended;
+			this.endEmitter.fire();
+		}
+	}
 
 	private ensureTestIsKnown(test: vscode.TestItem) {
 		if (!(test instanceof TestItemImpl)) {
@@ -539,7 +565,7 @@ export class TestRunCoordinator {
 	public cancelRunById(runId: string) {
 		for (const tracker of this.tracked.values()) {
 			if (tracker.id === runId) {
-				tracker.dispose();
+				tracker.cancel();
 				return;
 			}
 		}
@@ -550,7 +576,7 @@ export class TestRunCoordinator {
 	 */
 	public cancelAllRuns() {
 		for (const tracker of this.tracked.values()) {
-			tracker.dispose();
+			tracker.cancel();
 		}
 	}
 
@@ -578,7 +604,11 @@ export class TestRunCoordinator {
 		});
 
 		const tracker = this.getTracker(request, dto);
-		tracker.onEnd(() => this.proxy.$finishedExtensionTestRun(dto.id));
+		tracker.onEnd(() => {
+			this.proxy.$finishedExtensionTestRun(dto.id);
+			tracker.dispose();
+		});
+
 		return tracker.createRun(name);
 	}
 

src/vs/workbench/api/test/browser/extHostTesting.test.ts

@@ -4,21 +4,22 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as assert from 'assert';
+import * as sinon from 'sinon';
 import { VSBuffer } from 'vs/base/common/buffer';
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
 import { Iterable } from 'vs/base/common/iterator';
 import { URI } from 'vs/base/common/uri';
 import { mockObject, MockObject } from 'vs/base/test/common/mock';
+import * as editorRange from 'vs/editor/common/core/range';
 import { MainThreadTestingShape } from 'vs/workbench/api/common/extHost.protocol';
+import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 import { TestRunCoordinator, TestRunDto, TestRunProfileImpl } from 'vs/workbench/api/common/extHostTesting';
 import { ExtHostTestItemCollection, TestItemImpl } from 'vs/workbench/api/common/extHostTestItem';
 import * as convert from 'vs/workbench/api/common/extHostTypeConverters';
 import { Location, Position, Range, TestMessage, TestResultState, TestRunProfileKind, TestRunRequest as TestRunRequestImpl, TestTag } from 'vs/workbench/api/common/extHostTypes';
-import { TestDiffOpType, TestItemExpandState, TestMessageType, TestsDiff } from 'vs/workbench/contrib/testing/common/testTypes';
 import { TestId } from 'vs/workbench/contrib/testing/common/testId';
+import { TestDiffOpType, TestItemExpandState, TestMessageType, TestsDiff } from 'vs/workbench/contrib/testing/common/testTypes';
 import type { TestItem, TestRunRequest } from 'vscode';
-import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
-import * as editorRange from 'vs/editor/common/core/range';
 
 const simplify = (item: TestItem) => ({
 	id: item.id,
@@ -620,12 +621,12 @@ suite('ExtHost Testing', () => {
 
 		test('tracks a run started from a main thread request', () => {
 			const tracker = c.prepareForMainThreadTestRun(req, dto, cts.token);
-			assert.strictEqual(tracker.isRunning, false);
+			assert.strictEqual(tracker.hasRunningTasks, false);
 
 			const task1 = c.createTestRun('ctrl', single, req, 'run1', true);
 			const task2 = c.createTestRun('ctrl', single, req, 'run2', true);
 			assert.strictEqual(proxy.$startedExtensionTestRun.called, false);
-			assert.strictEqual(tracker.isRunning, true);
+			assert.strictEqual(tracker.hasRunningTasks, true);
 
 			task1.appendOutput('hello');
 			const taskId = proxy.$appendOutputToRun.args[0]?.[1];
@@ -633,19 +634,65 @@ suite('ExtHost Testing', () => {
 			task1.end();
 
 			assert.strictEqual(proxy.$finishedExtensionTestRun.called, false);
-			assert.strictEqual(tracker.isRunning, true);
+			assert.strictEqual(tracker.hasRunningTasks, true);
 
 			task2.end();
 
 			assert.strictEqual(proxy.$finishedExtensionTestRun.called, false);
-			assert.strictEqual(tracker.isRunning, false);
+			assert.strictEqual(tracker.hasRunningTasks, false);
+		});
+
+		test('run cancel force ends after a timeout', () => {
+			const clock = sinon.useFakeTimers();
+			try {
+				const tracker = c.prepareForMainThreadTestRun(req, dto, cts.token);
+				const task = c.createTestRun('ctrl', single, req, 'run1', true);
+				const onEnded = sinon.stub();
+				tracker.onEnd(onEnded);
+
+				assert.strictEqual(task.token.isCancellationRequested, false);
+				assert.strictEqual(tracker.hasRunningTasks, true);
+				tracker.cancel();
+
+				assert.strictEqual(task.token.isCancellationRequested, true);
+				assert.strictEqual(tracker.hasRunningTasks, true);
+
+				clock.tick(9999);
+				assert.strictEqual(tracker.hasRunningTasks, true);
+				assert.strictEqual(onEnded.called, false);
+
+				clock.tick(1);
+				assert.strictEqual(onEnded.called, true);
+				assert.strictEqual(tracker.hasRunningTasks, false);
+			} finally {
+				clock.restore();
+			}
+		});
+
+		test('run cancel force ends on second cancellation request', () => {
+			const tracker = c.prepareForMainThreadTestRun(req, dto, cts.token);
+			const task = c.createTestRun('ctrl', single, req, 'run1', true);
+			const onEnded = sinon.stub();
+			tracker.onEnd(onEnded);
+
+			assert.strictEqual(task.token.isCancellationRequested, false);
+			assert.strictEqual(tracker.hasRunningTasks, true);
+			tracker.cancel();
+
+			assert.strictEqual(task.token.isCancellationRequested, true);
+			assert.strictEqual(tracker.hasRunningTasks, true);
+			assert.strictEqual(onEnded.called, false);
+			tracker.cancel();
+
+			assert.strictEqual(tracker.hasRunningTasks, false);
+			assert.strictEqual(onEnded.called, true);
 		});
 
 		test('tracks a run started from an extension request', () => {
 			const task1 = c.createTestRun('ctrl', single, req, 'hello world', false);
 
 			const tracker = Iterable.first(c.trackers)!;
-			assert.strictEqual(tracker.isRunning, true);
+			assert.strictEqual(tracker.hasRunningTasks, true);
 			assert.deepStrictEqual(proxy.$startedExtensionTestRun.args, [
 				[{
 					profile: { group: 2, id: 42 },
@@ -662,11 +709,11 @@ suite('ExtHost Testing', () => {
 
 			task1.end();
 			assert.strictEqual(proxy.$finishedExtensionTestRun.called, false);
-			assert.strictEqual(tracker.isRunning, true);
+			assert.strictEqual(tracker.hasRunningTasks, true);
 
 			task2.end();
 			assert.deepStrictEqual(proxy.$finishedExtensionTestRun.args, [[tracker.id]]);
-			assert.strictEqual(tracker.isRunning, false);
+			assert.strictEqual(tracker.hasRunningTasks, false);
 
 			task3Detached.end();
 		});

src/vs/workbench/contrib/testing/browser/testingProgressUiService.ts

@@ -94,7 +94,7 @@ export class TestingProgressUiService extends Disposable implements ITestingProg
 	private readonly testViewProg = this._register(new MutableDisposable<UnmanagedProgress>());
 	private readonly updateCountsEmitter = new Emitter<CountSummary>();
 	private readonly updateTextEmitter = new Emitter<string>();
-	private lastRunSoFar = 0;
+	private lastProgress = 0;
 
 	public readonly onCountChange = this.updateCountsEmitter.event;
 	public readonly onTextChange = this.updateTextEmitter.event;
@@ -122,7 +122,7 @@ export class TestingProgressUiService extends Disposable implements ITestingProg
 
 			this.windowProg.clear();
 			this.testViewProg.clear();
-			this.lastRunSoFar = 0;
+			this.lastProgress = 0;
 			return;
 		}
 
@@ -133,7 +133,7 @@ export class TestingProgressUiService extends Disposable implements ITestingProg
 			});
 			this.testViewProg.value = this.instantiaionService.createInstance(UnmanagedProgress, {
 				location: Testing.ViewletId,
-				total: 100,
+				total: 1000,
 			});
 		}
 
@@ -143,8 +143,10 @@ export class TestingProgressUiService extends Disposable implements ITestingProg
 		const message = getTestProgressText(true, collected);
 		this.updateTextEmitter.fire(message);
 		this.windowProg.value.report({ message });
-		this.testViewProg.value!.report({ increment: collected.runSoFar - this.lastRunSoFar, total: collected.totalWillBeRun });
-		this.lastRunSoFar = collected.runSoFar;
+		const nextProgress = collected.runSoFar / collected.totalWillBeRun;
+		console.log({ increment: nextProgress - this.lastProgress, total: 1 });
+		this.testViewProg.value!.report({ increment: (nextProgress - this.lastProgress) * 1000, total: 1 });
+		this.lastProgress = nextProgress;
 	}
 }