Fix Notebook toolbar contributed actions (#11572)

* Fix Notebook toolbar contributed actions

* Fix args

* Remove unused imports

extensions/notebook/src/jupyter/jupyterController.ts

@@ -201,9 +201,9 @@ export class JupyterController implements vscode.Disposable {
 		});
 	}
 
-	public async doManagePackages(options?: ManagePackageDialogOptions): Promise<void> {
+	public async doManagePackages(options?: ManagePackageDialogOptions | vscode.Uri): Promise<void> {
 		try {
-			if (!options) {
+			if (!options || options instanceof vscode.Uri) {
 				options = {
 					defaultLocation: constants.localhostName,
 					defaultProviderId: LocalPipPackageManageProvider.ProviderId

src/sql/workbench/contrib/notebook/browser/cellViews/cellToolbar.component.ts

@@ -56,12 +56,12 @@ export class CellToolbarComponent {
 		this._actionBar = new Taskbar(taskbar);
 		this._actionBar.context = context;
 
-		let addCellsButton = new AddCellAction('notebook.AddCodeCell', localize('codeCellsPreview', "Add cell"), 'notebook-button masked-pseudo code');
+		let addCellsButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddCodeCell', localize('codeCellsPreview', "Add cell"), 'notebook-button masked-pseudo code');
 
-		let addCodeCellButton = new AddCellAction('notebook.AddCodeCell', localize('codePreview', "Code cell"), 'notebook-button masked-pseudo code');
+		let addCodeCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddCodeCell', localize('codePreview', "Code cell"), 'notebook-button masked-pseudo code');
 		addCodeCellButton.cellType = CellTypes.Code;
 
-		let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('textPreview', "Text cell"), 'notebook-button masked-pseudo markdown');
+		let addTextCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddTextCell', localize('textPreview', "Text cell"), 'notebook-button masked-pseudo markdown');
 		addTextCellButton.cellType = CellTypes.Markdown;
 
 		let deleteButton = this.instantiationService.createInstance(DeleteCellAction, 'delete', 'codicon masked-icon delete', localize('delete', "Delete"));

src/sql/workbench/contrib/notebook/browser/notebook.component.ts

@@ -322,7 +322,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		this._model = this._register(model);
 		await this._model.loadContents(trusted);
 		this.setLoading(false);
-		this.updateToolbarComponents(this._model.trustedMode);
+		this.updateToolbarComponents();
 		this.detectChanges();
 	}
 
@@ -348,11 +348,10 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		}
 	}
 
-	// Updates toolbar components
+	private updateToolbarComponents() {
-	private updateToolbarComponents(isTrusted: boolean) {
+		if (this.model.trustedMode) {
-		if (this._trustedAction) {
 			this._trustedAction.enabled = true;
-			this._trustedAction.trusted = isTrusted;
+			this._trustedAction.trusted = true;
 		}
 	}
 
@@ -417,26 +416,26 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			let spacerElement = document.createElement('li');
 			spacerElement.style.marginLeft = 'auto';
 
-			let addCellsButton = new AddCellAction('notebook.AddCodeCell', localize('codeCellsPreview', "Add cell"), 'notebook-button masked-pseudo code');
+			let addCellsButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddCodeCell', localize('codeCellsPreview', "Add cell"), 'notebook-button masked-pseudo code');
 
-			let addCodeCellButton = new AddCellAction('notebook.AddCodeCell', localize('codePreview', "Code cell"), 'notebook-button masked-pseudo code');
+			let addCodeCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddCodeCell', localize('codePreview', "Code cell"), 'notebook-button masked-pseudo code');
 			addCodeCellButton.cellType = CellTypes.Code;
 
-			let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('textPreview', "Text cell"), 'notebook-button masked-pseudo markdown');
+			let addTextCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddTextCell', localize('textPreview', "Text cell"), 'notebook-button masked-pseudo markdown');
 			addTextCellButton.cellType = CellTypes.Markdown;
 
 			this._runAllCellsAction = this.instantiationService.createInstance(RunAllCellsAction, 'notebook.runAllCells', localize('runAllPreview', "Run all"), 'notebook-button masked-pseudo start-outline');
 
 			let collapseCellsAction = this.instantiationService.createInstance(CollapseCellsAction, 'notebook.collapseCells', true);
 
-			let clearResultsButton = new ClearAllOutputsAction('notebook.ClearAllOutputs', true);
+			let clearResultsButton = this.instantiationService.createInstance(ClearAllOutputsAction, 'notebook.ClearAllOutputs', true);
 
 			this._trustedAction = this.instantiationService.createInstance(TrustedAction, 'notebook.Trusted', true);
 			this._trustedAction.enabled = false;
 
 			let taskbar = <HTMLElement>this.toolbar.nativeElement;
 			this._actionBar = new Taskbar(taskbar, { actionViewItemProvider: action => this.actionItemProvider(action as Action) });
-			this._actionBar.context = this;
+			this._actionBar.context = this._notebookParams.notebookUri;
 			taskbar.classList.add('in-preview');
 
 			let buttonDropdownContainer = DOM.$('li.action-item');
@@ -453,7 +452,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 				undefined
 			);
 			dropdownMenuActionViewItem.render(buttonDropdownContainer);
-			dropdownMenuActionViewItem.setActionContext(this);
+			dropdownMenuActionViewItem.setActionContext(this._notebookParams.notebookUri);
 
 			this._actionBar.setContent([
 				{ element: buttonDropdownContainer },
@@ -478,15 +477,15 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			attachToDropdown.render(attachToContainer);
 			attachSelectBoxStyler(attachToDropdown, this.themeService);
 
-			let addCodeCellButton = new AddCellAction('notebook.AddCodeCell', localize('code', "Code"), 'notebook-button icon-add');
+			let addCodeCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddCodeCell', localize('code', "Code"), 'notebook-button icon-add');
 			addCodeCellButton.cellType = CellTypes.Code;
 
-			let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('text', "Text"), 'notebook-button icon-add');
+			let addTextCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddTextCell', localize('text', "Text"), 'notebook-button icon-add');
 			addTextCellButton.cellType = CellTypes.Markdown;
 
 			this._runAllCellsAction = this.instantiationService.createInstance(RunAllCellsAction, 'notebook.runAllCells', localize('runAll', "Run Cells"), 'notebook-button icon-run-cells');
 
-			let clearResultsButton = new ClearAllOutputsAction('notebook.ClearAllOutputs', false);
+			let clearResultsButton = this.instantiationService.createInstance(ClearAllOutputsAction, 'notebook.ClearAllOutputs', false);
 
 			this._trustedAction = this.instantiationService.createInstance(TrustedAction, 'notebook.Trusted', false);
 			this._trustedAction.enabled = false;
@@ -495,7 +494,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 			let taskbar = <HTMLElement>this.toolbar.nativeElement;
 			this._actionBar = new Taskbar(taskbar, { actionViewItemProvider: action => this.actionItemProvider(action as Action) });
-			this._actionBar.context = this;
+			this._actionBar.context = this._notebookParams.notebookUri;
 
 			this._actionBar.setContent([
 				{ action: addCodeCellButton },

src/sql/workbench/contrib/notebook/browser/notebookActions.ts

@@ -25,10 +25,10 @@ import { INotebookModel } from 'sql/workbench/services/notebook/browser/models/m
 import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/browser/objectExplorerService';
 import { TreeUpdateUtils } from 'sql/workbench/services/objectExplorer/browser/treeUpdateUtils';
 import { find, firstIndex } from 'vs/base/common/arrays';
-import { INotebookEditor } from 'sql/workbench/services/notebook/browser/notebookService';
+import { INotebookService } from 'sql/workbench/services/notebook/browser/notebookService';
-import { NotebookComponent } from 'sql/workbench/contrib/notebook/browser/notebook.component';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { CellContext } from 'sql/workbench/contrib/notebook/browser/cellViews/codeActions';
+import { URI } from 'vs/base/common/uri';
 
 const msgLoading = localize('loading', "Loading kernels...");
 export const msgChanging = localize('changing', "Changing kernel...");
@@ -46,11 +46,12 @@ export class AddCellAction extends Action {
 	public cellType: CellType;
 
 	constructor(
-		id: string, label: string, cssClass: string
+		id: string, label: string, cssClass: string,
+		@INotebookService private _notebookService: INotebookService
 	) {
 		super(id, label, cssClass);
 	}
-	public async run(context: INotebookEditor | CellContext): Promise<void> {
+	public async run(context: URI | CellContext): Promise<void> {
 		let index = 0;
 		if (context instanceof CellContext) {
 			if (context?.model?.cells) {
@@ -64,14 +65,9 @@ export class AddCellAction extends Action {
 			}
 		} else {
 			//Add Cell after current selected cell.
-			if (context?.cells) {
+			const editor = this._notebookService.findNotebookEditor(context);
-				let notebookcomponent = context as NotebookComponent;
+			const index = editor.cells?.findIndex(cell => cell.active) ?? 0;
-				let id = notebookcomponent.activeCellId;
+			editor.addCell(this.cellType, index);
-				if (id) {
-					index = context.cells.findIndex(cell => cell.id === id) + 1;
-				}
-			}
-			context.addCell(this.cellType, index);
 		}
 	}
 }
@@ -110,7 +106,8 @@ export class ClearAllOutputsAction extends TooltipFromLabelAction {
 	private static readonly iconClass = 'icon-clear-results';
 	private static readonly maskedIconClass = 'masked-icon';
 
-	constructor(id: string, toggleTooltip: boolean) {
+	constructor(id: string, toggleTooltip: boolean,
+		@INotebookService private _notebookService: INotebookService) {
 		super(id, {
 			label: ClearAllOutputsAction.label,
 			baseClass: ClearAllOutputsAction.baseClass,
@@ -120,8 +117,9 @@ export class ClearAllOutputsAction extends TooltipFromLabelAction {
 		});
 	}
 
-	public run(context: INotebookEditor): Promise<boolean> {
+	public run(context: URI): Promise<boolean> {
-		return context.clearAllOutputs();
+		const editor = this._notebookService.findNotebookEditor(context);
+		return editor.clearAllOutputs();
 	}
 }
 
@@ -181,7 +179,8 @@ export class TrustedAction extends ToggleableAction {
 	private static readonly maskedIconClass = 'masked-icon';
 
 	constructor(
-		id: string, toggleTooltip: boolean
+		id: string, toggleTooltip: boolean,
+		@INotebookService private _notebookService: INotebookService
 	) {
 		super(id, {
 			baseClass: TrustedAction.baseClass,
@@ -202,17 +201,11 @@ export class TrustedAction extends ToggleableAction {
 		this.toggle(value);
 	}
 
-	public run(context: INotebookEditor): Promise<boolean> {
+	public async run(context: URI): Promise<boolean> {
-		let self = this;
+		const editor = this._notebookService.findNotebookEditor(context);
-		return new Promise<boolean>((resolve, reject) => {
+		this.trusted = !this.trusted;
-			try {
+		editor.model.trustedMode = this.trusted;
-				self.trusted = !self.trusted;
+		return true;
-				context.model.trustedMode = self.trusted;
-				resolve(true);
-			} catch (e) {
-				reject(e);
-			}
-		});
 	}
 }
 
@@ -220,13 +213,15 @@ export class TrustedAction extends ToggleableAction {
 export class RunAllCellsAction extends Action {
 	constructor(
 		id: string, label: string, cssClass: string,
-		@INotificationService private notificationService: INotificationService
+		@INotificationService private notificationService: INotificationService,
+		@INotebookService private _notebookService: INotebookService
 	) {
 		super(id, label, cssClass);
 	}
-	public async run(context: INotebookEditor): Promise<boolean> {
+	public async run(context: URI): Promise<boolean> {
 		try {
-			await context.runAllCells();
+			const editor = this._notebookService.findNotebookEditor(context);
+			await editor.runAllCells();
 			return true;
 		} catch (e) {
 			this.notificationService.error(getErrorMessage(e));
@@ -245,7 +240,8 @@ export class CollapseCellsAction extends ToggleableAction {
 	private static readonly expandCssClass = 'icon-show-cells';
 	private static readonly maskedIconClass = 'masked-icon';
 
-	constructor(id: string, toggleTooltip: boolean) {
+	constructor(id: string, toggleTooltip: boolean,
+		@INotebookService private _notebookService: INotebookService) {
 		super(id, {
 			baseClass: CollapseCellsAction.baseClass,
 			toggleOnLabel: CollapseCellsAction.expandCells,
@@ -265,19 +261,13 @@ export class CollapseCellsAction extends ToggleableAction {
 		this.toggle(value);
 	}
 
-	public run(context: INotebookEditor): Promise<boolean> {
+	public async run(context: URI): Promise<boolean> {
-		let self = this;
+		const editor = this._notebookService.findNotebookEditor(context);
-		return new Promise<boolean>((resolve, reject) => {
+		this.setCollapsed(!this.isCollapsed);
-			try {
+		editor.cells.forEach(cell => {
-				self.setCollapsed(!self.isCollapsed);
+			cell.isCollapsed = this.isCollapsed;
-				context.cells.forEach(cell => {
-					cell.isCollapsed = self.isCollapsed;
-				});
-				resolve(true);
-			} catch (e) {
-				reject(e);
-			}
 		});
+		return true;
 	}
 }
 

src/sql/workbench/contrib/notebook/test/browser/notebookActions.test.ts

@@ -8,11 +8,11 @@ import * as azdata from 'azdata';
 import * as sinon from 'sinon';
 import { TestConfigurationService } from 'sql/platform/connection/test/common/testConfigurationService';
 import { AddCellAction, ClearAllOutputsAction, CollapseCellsAction, KernelsDropdown, msgChanging, NewNotebookAction, noKernelName, RunAllCellsAction, TrustedAction } from 'sql/workbench/contrib/notebook/browser/notebookActions';
-import { ClientSessionStub, ContextViewProviderStub, NotebookComponentStub, NotebookModelStub } from 'sql/workbench/contrib/notebook/test/stubs';
+import { ClientSessionStub, ContextViewProviderStub, NotebookComponentStub, NotebookModelStub, NotebookServiceStub } from 'sql/workbench/contrib/notebook/test/stubs';
 import { NotebookEditorStub } from 'sql/workbench/contrib/notebook/test/testCommon';
 import { ICellModel, INotebookModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
 import { IStandardKernelWithProvider } from 'sql/workbench/services/notebook/browser/models/notebookUtils';
-import { INotebookEditor } from 'sql/workbench/services/notebook/browser/notebookService';
+import { INotebookEditor, INotebookService } from 'sql/workbench/services/notebook/browser/notebookService';
 import { CellType } from 'sql/workbench/services/notebook/common/contracts';
 import * as TypeMoq from 'typemoq';
 import { Emitter } from 'vs/base/common/event';
@@ -23,6 +23,7 @@ import { TestInstantiationService } from 'vs/platform/instantiation/test/common/
 import { INotificationService } from 'vs/platform/notification/common/notification';
 import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
 import { workbenchInstantiationService } from 'vs/workbench/test/browser/workbenchTestServices';
+import { URI } from 'vs/base/common/uri';
 
 class TestClientSession extends ClientSessionStub {
 	private _errorState: boolean = false;
@@ -103,122 +104,139 @@ class TestNotebookModel extends NotebookModelStub {
 }
 
 suite('Notebook Actions', function (): void {
+
+	let mockNotebookEditor: TypeMoq.Mock<INotebookEditor>;
+	let mockNotebookService: TypeMoq.Mock<INotebookService>;
+	const testUri = URI.parse('untitled');
+
+	suiteSetup(function (): void {
+		mockNotebookEditor = TypeMoq.Mock.ofType<INotebookEditor>(NotebookEditorStub);
+		mockNotebookService = TypeMoq.Mock.ofType<INotebookService>(NotebookServiceStub);
+		mockNotebookService.setup(x => x.findNotebookEditor(TypeMoq.It.isAny())).returns(uri => mockNotebookEditor.object);
+	});
+
+	teardown(function (): void {
+		mockNotebookEditor.reset();
+	});
+
 	test('Add Cell Action', async function (): Promise<void> {
 		let testCellType: CellType = 'code';
 		let actualCellType: CellType;
 
-		let action = new AddCellAction('TestId', 'TestLabel', 'TestClass');
+
+		let action = new AddCellAction('TestId', 'TestLabel', 'TestClass', mockNotebookService.object);
 		action.cellType = testCellType;
 
 		// Normal use case
+		mockNotebookEditor.setup(x => x.addCell(TypeMoq.It.isAny(), TypeMoq.It.isAnyNumber())).returns((cellType, index) => { actualCellType = cellType; });
 		let mockNotebookComponent = TypeMoq.Mock.ofType<INotebookEditor>(NotebookComponentStub);
 		mockNotebookComponent.setup(c => c.addCell(TypeMoq.It.isAny(), TypeMoq.It.isAnyNumber())).returns(cellType => {
 			actualCellType = cellType;
 		});
 
-		assert.doesNotThrow(() => action.run(mockNotebookComponent.object));
+		assert.doesNotThrow(() => action.run(testUri));
 		assert.strictEqual(actualCellType, testCellType);
 
 		// Handle error case
-		mockNotebookComponent.reset();
+		mockNotebookEditor.reset();
-		mockNotebookComponent.setup(c => c.addCell(TypeMoq.It.isAny(), TypeMoq.It.isAnyNumber())).throws(new Error('Test Error'));
+		mockNotebookEditor.setup(x => x.addCell(TypeMoq.It.isAny(), TypeMoq.It.isAnyNumber())).throws(new Error('Test Error'));
-		await assert.rejects(action.run(mockNotebookComponent.object));
+		await assert.rejects(action.run(URI.parse('untitled')));
 	});
 
 	test('Clear All Outputs Action', async function (): Promise<void> {
-		let action = new ClearAllOutputsAction('TestId', true);
+		let action = new ClearAllOutputsAction('TestId', true, mockNotebookService.object);
 
 		// Normal use case
-		let mockNotebookComponent = TypeMoq.Mock.ofType<INotebookEditor>(NotebookComponentStub);
+		mockNotebookEditor.setup(c => c.clearAllOutputs()).returns(() => Promise.resolve(true));
-		mockNotebookComponent.setup(c => c.clearAllOutputs()).returns(() => Promise.resolve(true));
 
-		let result = await action.run(mockNotebookComponent.object);
+		let result = await action.run(testUri);
 		assert.ok(result, 'Clear All Outputs Action should succeed');
-		mockNotebookComponent.verify(c => c.clearAllOutputs(), TypeMoq.Times.once());
+		mockNotebookEditor.verify(c => c.clearAllOutputs(), TypeMoq.Times.once());
 
 		// Handle failure case
-		mockNotebookComponent.reset();
+		mockNotebookEditor.reset();
-		mockNotebookComponent.setup(c => c.clearAllOutputs()).returns(() => Promise.resolve(false));
+		mockNotebookEditor.setup(c => c.clearAllOutputs()).returns(() => Promise.resolve(false));
 
-		result = await action.run(mockNotebookComponent.object);
+		result = await action.run(testUri);
 		assert.strictEqual(result, false, 'Clear All Outputs Action should have failed');
-		mockNotebookComponent.verify(c => c.clearAllOutputs(), TypeMoq.Times.once());
+		mockNotebookEditor.verify(c => c.clearAllOutputs(), TypeMoq.Times.once());
 	});
 
 	test('Trusted Action', async function (): Promise<void> {
 		let mockNotification = TypeMoq.Mock.ofType<INotificationService>(TestNotificationService);
 		mockNotification.setup(n => n.notify(TypeMoq.It.isAny()));
 
-		let action = new TrustedAction('TestId', true);
+		let action = new TrustedAction('TestId', true, mockNotebookService.object);
 		assert.strictEqual(action.trusted, false, 'Should not be trusted by default');
 
-		// Normal use case
+		const testNotebookModel: INotebookModel = <INotebookModel>{
-		let contextStub = <INotebookEditor>{
+			trustedMode: false
-			model: <INotebookModel>{
-				trustedMode: false
-			}
 		};
-		let result = await action.run(contextStub);
+
+		mockNotebookEditor.setup(x => x.model).returns(() => testNotebookModel);
+		// Normal use case
+		let result = await action.run(testUri);
 		assert.ok(result, 'Trusted Action should succeed');
 		assert.strictEqual(action.trusted, true, 'Should be trusted after toggling trusted state');
+		assert.strictEqual(testNotebookModel.trustedMode, true, 'Model should be true after toggling trusted state');
 
 		// Should toggle trusted to false on subsequent action
-		result = await action.run(contextStub);
+		result = await action.run(testUri);
 		assert.ok(result, 'Trusted Action should succeed again');
 		assert.strictEqual(action.trusted, false, 'Should toggle trusted to false');
+		assert.strictEqual(testNotebookModel.trustedMode, false, 'Model should be false again after toggling trusted state');
 	});
 
 	test('Run All Cells Action', async function (): Promise<void> {
 		let mockNotification = TypeMoq.Mock.ofType<INotificationService>(TestNotificationService);
 		mockNotification.setup(n => n.notify(TypeMoq.It.isAny()));
 
-		let action = new RunAllCellsAction('TestId', 'TestLabel', 'TestClass', mockNotification.object);
+		let action = new RunAllCellsAction('TestId', 'TestLabel', 'TestClass', mockNotification.object, mockNotebookService.object);
 
 		// Normal use case
-		let mockNotebookComponent = TypeMoq.Mock.ofType<INotebookEditor>(NotebookComponentStub);
+		mockNotebookEditor.setup(c => c.runAllCells()).returns(() => Promise.resolve(true));
-		mockNotebookComponent.setup(c => c.runAllCells()).returns(() => Promise.resolve(true));
 
-		let result = await action.run(mockNotebookComponent.object);
+		let result = await action.run(testUri);
 		assert.ok(result, 'Run All Cells Action should succeed');
-		mockNotebookComponent.verify(c => c.runAllCells(), TypeMoq.Times.once());
+		mockNotebookEditor.verify(c => c.runAllCells(), TypeMoq.Times.once());
 
 		// Handle errors
-		mockNotebookComponent.reset();
+		mockNotebookEditor.reset();
-		mockNotebookComponent.setup(c => c.runAllCells()).returns(() => { throw new Error('Test Error'); });
+		mockNotebookEditor.setup(c => c.runAllCells()).throws(new Error('Test Error'));
 
-		result = await action.run(mockNotebookComponent.object);
+		result = await action.run(testUri);
 		assert.strictEqual(result, false, 'Run All Cells Action should fail on error');
 	});
 
 	test('Collapse Cells Action', async function (): Promise<void> {
-		let action = new CollapseCellsAction('TestId', true);
+		let action = new CollapseCellsAction('TestId', true, mockNotebookService.object);
 		assert.strictEqual(action.isCollapsed, false, 'Should not be collapsed by default');
 
-		let context = <INotebookEditor>{
+		const testCells = [<ICellModel>{
-			cells: [<ICellModel>{
+			isCollapsed: false
-				isCollapsed: false
+		}, <ICellModel>{
-			}, <ICellModel>{
+			isCollapsed: true
-				isCollapsed: true
+		}, <ICellModel>{
-			}, <ICellModel>{
+			isCollapsed: false
-				isCollapsed: false
+		}];
-			}]
+
-		};
+		mockNotebookEditor.setup(x => x.cells).returns(() => testCells);
 
 		// Collapse cells case
-		let result = await action.run(context);
+		let result = await action.run(testUri);
 		assert.ok(result, 'Collapse Cells Action should succeed');
 
 		assert.strictEqual(action.isCollapsed, true, 'Action should be collapsed after first toggle');
-		context.cells.forEach(cell => {
+		testCells.forEach(cell => {
 			assert.strictEqual(cell.isCollapsed, true, 'Cells should be collapsed after first toggle');
 		});
 
 		// Toggle cells to uncollapsed
-		result = await action.run(context);
+		result = await action.run(testUri);
 		assert.ok(result, 'Collapse Cells Action should succeed');
 
 		assert.strictEqual(action.isCollapsed, false, 'Action should not be collapsed after second toggle');
-		context.cells.forEach(cell => {
+		testCells.forEach(cell => {
 			assert.strictEqual(cell.isCollapsed, false, 'Cells should not be collapsed after second toggle');
 		});
 	});
@@ -239,7 +257,6 @@ suite('Notebook Actions', function (): void {
 		assert.strictEqual(actualCmdId, NewNotebookAction.INTERNAL_NEW_NOTEBOOK_CMD_ID);
 	});
 
-
 	suite('Kernels dropdown', async () => {
 		let kernelsDropdown: KernelsDropdown;
 		let contextViewProvider: ContextViewProviderStub;

src/sql/workbench/contrib/notebook/test/testCommon.ts

@@ -5,7 +5,7 @@
 
 import { QueryTextEditor } from 'sql/workbench/browser/modelComponents/queryTextEditor';
 import * as stubs from 'sql/workbench/contrib/notebook/test/stubs';
-import { INotebookModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
+import { INotebookModel, ICellModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
 import { INotebookParams } from 'sql/workbench/services/notebook/browser/notebookService';
 import * as dom from 'vs/base/browser/dom';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
@@ -19,6 +19,7 @@ import { TestStorageService } from 'vs/workbench/test/common/workbenchTestServic
 export class NotebookEditorStub extends stubs.NotebookEditorStub {
 	cellEditors: CellEditorProviderStub[];
 	model: INotebookModel | undefined;
+	cells?: ICellModel[] = [];
 
 	get id(): string {
 		return this.notebookParams?.notebookUri?.toString();