Plots service unit tests (#4628)

Add some unit tests for the plots service Fix disposing and ensuring plot is removed Signed-off-by: Tim Mok <[email protected]> Co-authored-by: Wasim Lorgat <[email protected]>
posit-dev · Sep 16, 2024 · 57b320b · 57b320b
1 parent 02d8da3
commit 57b320b
Show file tree

Hide file tree

Showing 5 changed files with 264 additions and 20 deletions.
diff --git a/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts b/src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts
@@ -191,7 +191,7 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 
 		// When the storage service is about to save state, store the current history policy
 		// and storage policy in the workspace storage.
-		this._storageService.onWillSaveState(() => {
+		this._register(this._storageService.onWillSaveState(() => {
 
 			this._storageService.store(
 				HistoryPolicyStorageKey,
@@ -220,7 +220,7 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 					StorageScope.WORKSPACE,
 					StorageTarget.MACHINE);
 			}
-		});
+		}));
 
 		// When the extension service is about to stop, remove any HTML plots
 		// from the plots list. These plots are backed by a proxy that runs in
@@ -770,10 +770,10 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 	 */
 	removePlot(id: string): void {
 		// Find the plot with the given ID and dispose it
-		// It will be automatically removed from the list during onDidClose
 		this._plots.forEach((plot, index) => {
 			if (plot.id === id) {
 				this.unregisterPlotClient(plot);
+				this._plots.splice(index, 1);
 			}
 		});
 
@@ -1095,6 +1095,8 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 			this._plotClientsByComm.delete(metadata.id);
 		}));
 
+		this._register(commProxy);
+
 		return commProxy;
 	}
 

diff --git a/src/vs/workbench/contrib/positronPlots/test/browser/positronPlotsService.test.ts b/src/vs/workbench/contrib/positronPlots/test/browser/positronPlotsService.test.ts
@@ -0,0 +1,230 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (C) 2024 Posit Software, PBC. All rights reserved.
+ *  Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import assert from 'assert';
+import { raceTimeout, timeout } from 'vs/base/common/async';
+import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils';
+import { PositronIPyWidgetsService } from 'vs/workbench/contrib/positronIPyWidgets/browser/positronIPyWidgetsService';
+import { PositronPlotsService } from 'vs/workbench/contrib/positronPlots/browser/positronPlotsService';
+import { PositronWebviewPreloadService } from 'vs/workbench/contrib/positronWebviewPreloads/browser/positronWebviewPreloadsService';
+import { IPositronPlotMetadata } from 'vs/workbench/services/languageRuntime/common/languageRuntimePlotClient';
+import { LanguageRuntimeSessionMode } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService';
+import { IPositronIPyWidgetsService } from 'vs/workbench/services/positronIPyWidgets/common/positronIPyWidgetsService';
+import { HistoryPolicy, IPositronPlotClient } from 'vs/workbench/services/positronPlots/common/positronPlots';
+import { IPositronWebviewPreloadService } from 'vs/workbench/services/positronWebviewPreloads/common/positronWebviewPreloadService';
+import { IRuntimeSessionService, RuntimeClientType } from 'vs/workbench/services/runtimeSession/common/runtimeSessionService';
+import { TestLanguageRuntimeSession } from 'vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession';
+import { TestRuntimeSessionService } from 'vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService';
+import { IViewsService } from 'vs/workbench/services/views/common/viewsService';
+import { TestViewsService, workbenchInstantiationService } from 'vs/workbench/test/browser/workbenchTestServices';
+
+suite('Positron - Plots Service', () => {
+
+	const disposables = ensureNoDisposablesAreLeakedInTestSuite();
+	let plotsService: PositronPlotsService;
+	let runtimeSessionService: TestRuntimeSessionService;
+
+	setup(() => {
+		const instantiationService = workbenchInstantiationService(undefined, disposables);
+		runtimeSessionService = disposables.add(instantiationService.createInstance(TestRuntimeSessionService));
+		instantiationService.stub(IRuntimeSessionService, runtimeSessionService);
+		instantiationService.stub(IPositronWebviewPreloadService, disposables.add(instantiationService.createInstance(PositronWebviewPreloadService)));
+		instantiationService.stub(IPositronIPyWidgetsService, disposables.add(instantiationService.createInstance(PositronIPyWidgetsService)));
+		instantiationService.stub(IViewsService, new TestViewsService());
+
+		plotsService = disposables.add(instantiationService.createInstance(PositronPlotsService));
+	});
+
+	async function createSession() {
+		const session = disposables.add(new TestLanguageRuntimeSession(LanguageRuntimeSessionMode.Console));
+		runtimeSessionService.startSession(session);
+
+		await timeout(0);
+
+		const out: {
+			session: TestLanguageRuntimeSession;
+			plotClient: IPositronPlotClient | undefined;
+		} = {
+			session, plotClient: undefined,
+		};
+
+		disposables.add(session.onDidCreateClientInstance(client => out.plotClient = {
+			id: client.client.getClientId(),
+			metadata: {} as IPositronPlotMetadata,
+		} as IPositronPlotClient));
+
+		return out;
+	}
+
+	test('history policy: change history policy', () => {
+		plotsService.selectHistoryPolicy(HistoryPolicy.AlwaysVisible);
+		assert.strictEqual(plotsService.historyPolicy, HistoryPolicy.AlwaysVisible);
+
+		plotsService.selectHistoryPolicy(HistoryPolicy.Automatic);
+		assert.strictEqual(plotsService.historyPolicy, HistoryPolicy.Automatic);
+
+		plotsService.selectHistoryPolicy(HistoryPolicy.NeverVisible);
+		assert.strictEqual(plotsService.historyPolicy, HistoryPolicy.NeverVisible);
+	});
+
+	test('history policy: change event', async () => {
+		let historyPolicyChanged = 0;
+
+		const didChangeHistoryPolicy = new Promise<void>((resolve) => {
+			const disposable = plotsService.onDidChangeHistoryPolicy((e) => {
+				historyPolicyChanged++;
+				resolve();
+			});
+			disposables.add(disposable);
+		});
+
+		// no event since 'Automatic' is the default
+		plotsService.selectHistoryPolicy(HistoryPolicy.Automatic);
+
+		// event occurs when changing to 'AlwaysVisible'
+		plotsService.selectHistoryPolicy(HistoryPolicy.AlwaysVisible);
+
+		await raceTimeout(didChangeHistoryPolicy, 100, () => assert.fail('onDidChangeHistoryPolicy event did not fire'));
+		assert.strictEqual(historyPolicyChanged, 1, 'onDidChangeHistoryPolicy event should fire once');
+	});
+
+	test('sizing policy: check options and change size', () => {
+		assert.throws(() => plotsService.selectSizingPolicy('non-existant sizing policy'));
+
+		assert.strictEqual(plotsService.sizingPolicies.length, 6);
+
+		plotsService.selectSizingPolicy('auto');
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'auto');
+
+		plotsService.selectSizingPolicy('fill');
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'fill');
+
+		plotsService.selectSizingPolicy('landscape');
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'landscape');
+
+		plotsService.selectSizingPolicy('portrait');
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'portrait');
+
+		plotsService.selectSizingPolicy('square');
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'square');
+
+		plotsService.setCustomPlotSize({ width: 100, height: 100 });
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'custom');
+		assert.strictEqual(plotsService.sizingPolicies.length, 7);
+
+		plotsService.clearCustomPlotSize();
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'auto');
+		assert.strictEqual(plotsService.sizingPolicies.length, 6);
+	});
+
+	test('sizing policy: change event', async () => {
+		let sizingPolicyChanged = 0;
+
+		const didChangeSizingPolicy = new Promise<void>((resolve) => {
+			const disposable = plotsService.onDidChangeSizingPolicy((e) => {
+				sizingPolicyChanged++;
+				resolve();
+			});
+			disposables.add(disposable);
+		});
+
+		// no event since 'auto' is the default
+		plotsService.selectSizingPolicy('auto');
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'auto');
+
+		// event occurs when changing to 'fill'
+		plotsService.selectSizingPolicy('fill');
+		assert.strictEqual(plotsService.selectedSizingPolicy.id, 'fill');
+
+		await raceTimeout(didChangeSizingPolicy, 100, () => assert.fail('onDidChangeSizingPolicy event did not fire'));
+		assert.strictEqual(sizingPolicyChanged, 1, 'onDidChangeSizingPolicy event should fire once for changing to "fill"');
+	});
+
+	test('selection: select plot', async () => {
+		const session = await createSession();
+
+		session.session.createClient(RuntimeClientType.Plot, {}, {}, 'plot1');
+		session.session.createClient(RuntimeClientType.Plot, {}, {}, 'plot2');
+
+		assert.strictEqual(plotsService.selectedPlotId, 'plot2');
+
+		let selectPlotCalled = false;
+		const didSelectPlot = new Promise<void>((resolve) => {
+			const disposable = plotsService.onDidSelectPlot((e) => {
+				selectPlotCalled = true;
+				resolve();
+			});
+			disposables.add(disposable);
+		});
+		plotsService.selectPlot('plot1');
+
+		await raceTimeout(didSelectPlot, 100, () => assert.fail('onDidSelectPlot event did not fire'));
+
+		assert.ok(selectPlotCalled, 'onDidSelectPlot event should fire');
+		assert.strictEqual(plotsService.selectedPlotId, 'plot1');
+	});
+
+	test('selection: remove selected plot', async () => {
+		const session = await createSession();
+
+		session.session.createClient(RuntimeClientType.Plot, {}, {}, 'plot1');
+
+		let removePlotCalled = false;
+
+		const didRemovePlot = new Promise<void>((resolve) => {
+			const disposable = plotsService.onDidRemovePlot((e) => {
+				removePlotCalled = true;
+				resolve();
+			});
+			disposables.add(disposable);
+		});
+
+		assert.strictEqual(plotsService.selectedPlotId, 'plot1');
+
+		plotsService.removeSelectedPlot();
+
+		await raceTimeout(didRemovePlot, 100, () => assert.fail('onDidRemovePlot event did not fire'));
+
+		assert.ok(removePlotCalled, 'onDidRemovePlot event should fire');
+		assert.strictEqual(plotsService.positronPlotInstances.length, 0);
+		assert.strictEqual(plotsService.selectedPlotId, undefined);
+	});
+
+	test('selection: expect error removing plot when no plot selected', () => {
+		assert.throws(() => plotsService.removeSelectedPlot(), { message: 'No plot is selected' });
+	});
+
+	test('selection: select previous/next plot', async () => {
+		const session = await createSession();
+
+		session.session.createClient(RuntimeClientType.Plot, {}, {}, 'plot1');
+		session.session.createClient(RuntimeClientType.Plot, {}, {}, 'plot2');
+		session.session.createClient(RuntimeClientType.Plot, {}, {}, 'plot3');
+
+		assert.strictEqual(plotsService.selectedPlotId, 'plot3');
+
+		plotsService.selectPreviousPlot();
+		assert.strictEqual(plotsService.selectedPlotId, 'plot2');
+
+		plotsService.selectPreviousPlot();
+		assert.strictEqual(plotsService.selectedPlotId, 'plot1');
+
+		plotsService.selectNextPlot();
+		assert.strictEqual(plotsService.selectedPlotId, 'plot2');
+
+		plotsService.selectNextPlot();
+		assert.strictEqual(plotsService.selectedPlotId, 'plot3');
+	});
+
+	test('plot client: create client event', async () => {
+		const session = await createSession();
+
+		assert.strictEqual(plotsService.positronPlotInstances.length, 0);
+		session.session.createClient(RuntimeClientType.Plot, {}, {}, 'plot1');
+
+		assert.strictEqual(plotsService.selectedPlotId, 'plot1');
+		assert.strictEqual(plotsService.positronPlotInstances.length, 1);
+	});
+});
diff --git a/src/vs/workbench/services/languageRuntime/common/languageRuntimePlotClient.ts b/src/vs/workbench/services/languageRuntime/common/languageRuntimePlotClient.ts
@@ -143,13 +143,13 @@ export class PlotClientInstance extends Disposable implements IPositronPlotClien
 
 		// Connect close emitter event
 		this.onDidClose = this._closeEmitter.event;
-		this._commProxy.onDidClose((state) => {
+		this._register(this._commProxy.onDidClose((state) => {
 			this._closeEmitter.fire();
 
 			// Silently cancel any pending render requests
 			this._currentRender?.cancel();
 			this._stateEmitter.fire(PlotClientState.Closed);
-		});
+		}));
 
 		// Connect the state emitter event
 		this.onDidChangeState = this._stateEmitter.event;
@@ -167,20 +167,20 @@ export class PlotClientInstance extends Disposable implements IPositronPlotClien
 		this.onDidSetIntrinsicSize = this._didSetIntrinsicSizeEmitter.event;
 
 		// Listen to our own state changes
-		this.onDidChangeState((state) => {
+		this._register(this.onDidChangeState((state) => {
 			this._state = state;
-		});
+		}));
 
 		// Listen for plot updates
-		this._commProxy.onDidRenderUpdate(async () => {
+		this._register(this._commProxy.onDidRenderUpdate(async () => {
 			const rendered = await this.queuePlotUpdateRequest();
 			this._renderUpdateEmitter.fire(rendered);
-		});
+		}));
 
 		// Listn for plot show events
-		this._commProxy.onDidShowPlot(async (_evt) => {
+		this._register(this._commProxy.onDidShowPlot(async (_evt) => {
 			this._didShowPlotEmitter.fire();
-		});
+		}));
 	}
 
 	/**

diff --git a/src/vs/workbench/services/languageRuntime/common/positronPlotCommProxy.ts b/src/vs/workbench/services/languageRuntime/common/positronPlotCommProxy.ts
@@ -160,15 +160,15 @@ export class PositronPlotCommProxy extends Disposable {
 
 		// Connect close emitter event
 		this.onDidClose = this._closeEmitter.event;
-		clientStateEvent((state) => {
+		this._register(clientStateEvent((state) => {
 			if (state === RuntimeClientState.Closed) {
 				this._closeEmitter.fire();
 
 				// Silently cancel any pending render requests
 				this._currentRender?.cancel();
 				this._renderQueue.forEach((render) => render.cancel());
 			}
-		});
+		}));
 
 		// Connect the render update emitter event
 		this.onDidRenderUpdate = this._renderUpdateEmitter.event;
@@ -179,17 +179,17 @@ export class PositronPlotCommProxy extends Disposable {
 		// Connect the intrinsic size emitter event
 		this.onDidSetIntrinsicSize = this._didSetIntrinsicSizeEmitter.event;
 
-		this._comm.onDidClose(() => {
+		this._register(this._comm.onDidClose(() => {
 			this._closeEmitter.fire();
-		});
+		}));
 
-		this._comm.onDidShow(() => {
+		this._register(this._comm.onDidShow(() => {
 			this._didShowPlotEmitter.fire();
-		});
+		}));
 
-		this._comm.onDidUpdate((_evt) => {
+		this._register(this._comm.onDidUpdate((_evt) => {
 			this._renderUpdateEmitter.fire();
-		});
+		}));
 
 		this._register(this._comm);
 	}

diff --git a/src/vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService.ts b/src/vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService.ts
@@ -4,20 +4,32 @@
  *--------------------------------------------------------------------------------------------*/
 import { Emitter } from 'vs/base/common/event';
 import { Disposable } from 'vs/base/common/lifecycle';
-import { ILanguageRuntimeSession, IRuntimeSessionService, IRuntimeSessionWillStartEvent } from 'vs/workbench/services/runtimeSession/common/runtimeSessionService';
+import { ILanguageRuntimeClientCreatedEvent } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService';
+import { ILanguageRuntimeGlobalEvent, ILanguageRuntimeSession, IRuntimeSessionService, IRuntimeSessionWillStartEvent } from 'vs/workbench/services/runtimeSession/common/runtimeSessionService';
 
 export class TestRuntimeSessionService extends Disposable implements Partial<IRuntimeSessionService> {
 	private readonly _willStartEmitter = this._register(new Emitter<IRuntimeSessionWillStartEvent>());
+	private readonly _didStartRuntime = this._register(new Emitter<ILanguageRuntimeSession>());
+	private readonly _didReceiveRuntimeEvent = this._register(new Emitter<ILanguageRuntimeGlobalEvent>());
+	private readonly _didCreateClientInstance = this._register(new Emitter<ILanguageRuntimeClientCreatedEvent>());
 
 	readonly activeSessions = new Array<ILanguageRuntimeSession>();
 
 	readonly onWillStartSession = this._willStartEmitter.event;
 
+	readonly onDidStartRuntime = this._didStartRuntime.event;
+
+	readonly onDidReceiveRuntimeEvent = this._didReceiveRuntimeEvent.event;
+
+	readonly onDidCreateClientInstance = this._didCreateClientInstance.event;
+
 	// Test helpers.
 
 	startSession(session: ILanguageRuntimeSession): void {
 		this.activeSessions.push(session);
+		this._register(session.onDidCreateClientInstance(e => this._didCreateClientInstance.fire(e)));
 		this._willStartEmitter.fire({ session, isNew: true });
+		this._didStartRuntime.fire(session);
 	}
 }