Add units tests for Notebook Actions (#8595)

src/sql/workbench/contrib/notebook/browser/cellViews/codeActions.ts

@@ -12,13 +12,13 @@ import { INotificationService, Severity } from 'vs/platform/notification/common/
 import { NotebookModel } from 'sql/workbench/contrib/notebook/browser/models/notebookModel';
 import { ICellModel, CellExecutionState } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
-import { MultiStateAction, IMultiStateData } from 'sql/workbench/contrib/notebook/browser/notebookActions';
 import { IKeybindingService } from 'vs/platform/keybinding/common/keybinding';
 import { ILogService } from 'vs/platform/log/common/log';
 import { getErrorMessage } from 'vs/base/common/errors';
 
 let notebookMoreActionMsg = localize('notebook.failed', "Please select active cell and try again");
 const emptyExecutionCountLabel = '[ ]';
+const HIDE_ICON_CLASS = ' hideIcon';
 
 function hasModelAndCell(context: CellContext, notificationService: INotificationService): boolean {
 	if (!context || !context.model) {
@@ -63,6 +63,96 @@ export abstract class CellActionBase extends Action {
 	abstract doRun(context: CellContext): Promise<void>;
 }
 
+interface IActionStateData {
+	className?: string;
+	label?: string;
+	tooltip?: string;
+	hideIcon?: boolean;
+	commandId?: string;
+}
+
+class IMultiStateData<T> {
+	private _stateMap = new Map<T, IActionStateData>();
+	constructor(mappings: { key: T, value: IActionStateData }[], private _state: T, private _baseClass?: string) {
+		if (mappings) {
+			mappings.forEach(s => this._stateMap.set(s.key, s.value));
+		}
+	}
+
+	public set state(value: T) {
+		if (!this._stateMap.has(value)) {
+			throw new Error('State value must be in stateMap');
+		}
+		this._state = value;
+	}
+
+	public updateStateData(state: T, updater: (data: IActionStateData) => void): void {
+		let data = this._stateMap.get(state);
+		if (data) {
+			updater(data);
+		}
+	}
+
+	public get classes(): string {
+		let classVal = this.getStateValueOrDefault<string>((data) => data.className, '');
+		let classes = this._baseClass ? `${this._baseClass} ` : '';
+		classes += classVal;
+		if (this.getStateValueOrDefault<boolean>((data) => data.hideIcon, false)) {
+			classes += HIDE_ICON_CLASS;
+		}
+		return classes;
+	}
+
+	public get label(): string {
+		return this.getStateValueOrDefault<string>((data) => data.label, '');
+	}
+
+	public get tooltip(): string {
+		return this.getStateValueOrDefault<string>((data) => data.tooltip, '');
+	}
+
+	public get commandId(): string {
+		return this.getStateValueOrDefault<string>((data) => data.commandId, '');
+	}
+
+	private getStateValueOrDefault<U>(getter: (data: IActionStateData) => U, defaultVal?: U): U {
+		let data = this._stateMap.get(this._state);
+		return data ? getter(data) : defaultVal;
+	}
+}
+
+abstract class MultiStateAction<T> extends Action {
+	constructor(
+		id: string,
+		protected states: IMultiStateData<T>,
+		private _keybindingService: IKeybindingService,
+		private readonly logService: ILogService) {
+		super(id, '');
+		this.updateLabelAndIcon();
+	}
+
+	private updateLabelAndIcon() {
+		let keyboardShortcut: string;
+		try {
+			// If a keyboard shortcut exists for the command id passed in, append that to the label
+			if (this.states.commandId !== '') {
+				let binding = this._keybindingService.lookupKeybinding(this.states.commandId);
+				keyboardShortcut = binding ? binding.getLabel() : undefined;
+			}
+		} catch (error) {
+			this.logService.error(error);
+		}
+		this.label = this.states.label;
+		this.tooltip = keyboardShortcut ? this.states.tooltip + ` (${keyboardShortcut})` : this.states.tooltip;
+		this.class = this.states.classes;
+	}
+
+	protected updateState(state: T): void {
+		this.states.state = state;
+		this.updateLabelAndIcon();
+	}
+}
+
 export class RunCellAction extends MultiStateAction<CellExecutionState> {
 	public static ID = 'notebook.runCell';
 	public static LABEL = 'Run cell';

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

@@ -17,16 +17,14 @@ import { ConnectionProfile } from 'sql/platform/connection/common/connectionProf
 import { noKernel } from 'sql/workbench/services/notebook/browser/sessionManager';
 import { IConnectionDialogService } from 'sql/workbench/services/connection/common/connectionDialogService';
 import { NotebookModel } from 'sql/workbench/contrib/notebook/browser/models/notebookModel';
-import { IKeybindingService } from 'vs/platform/keybinding/common/keybinding';
-import { ILogService } from 'vs/platform/log/common/log';
 import { ICommandService } from 'vs/platform/commands/common/commands';
 import { CellType } from 'sql/workbench/contrib/notebook/common/models/contracts';
-import { NotebookComponent } from 'sql/workbench/contrib/notebook/browser/notebook.component';
 import { getErrorMessage } from 'vs/base/common/errors';
 import { INotebookModel } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
 import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/browser/objectExplorerService';
 import { TreeUpdateUtils } from 'sql/workbench/contrib/objectExplorer/browser/treeUpdateUtils';
 import { find, firstIndex } from 'vs/base/common/arrays';
+import { INotebookEditor } from 'sql/workbench/services/notebook/browser/notebookService';
 
 const msgLoading = localize('loading', "Loading kernels...");
 const msgChanging = localize('changing', "Changing kernel...");
@@ -36,7 +34,6 @@ const msgLoadingContexts = localize('loadingContexts', "Loading contexts...");
 const msgChangeConnection = localize('changeConnection', "Change Connection");
 const msgSelectConnection = localize('selectConnection', "Select Connection");
 const msgLocalHost = localize('localhost', "localhost");
-const HIDE_ICON_CLASS = ' hideIcon';
 
 // Action to add a cell to notebook based on cell type(code/markdown).
 export class AddCellAction extends Action {
@@ -47,7 +44,7 @@ export class AddCellAction extends Action {
 	) {
 		super(id, label, cssClass);
 	}
-	public run(context: NotebookComponent): Promise<boolean> {
+	public run(context: INotebookEditor): Promise<boolean> {
 		return new Promise<boolean>((resolve, reject) => {
 			try {
 				context.addCell(this.cellType);
@@ -59,7 +56,6 @@ export class AddCellAction extends Action {
 	}
 }
 
-
 // Action to clear outputs of all code cells.
 export class ClearAllOutputsAction extends Action {
 	constructor(
@@ -67,7 +63,7 @@ export class ClearAllOutputsAction extends Action {
 	) {
 		super(id, label, cssClass);
 	}
-	public run(context: NotebookComponent): Promise<boolean> {
+	public run(context: INotebookEditor): Promise<boolean> {
 		return context.clearAllOutputs();
 	}
 }
@@ -106,99 +102,6 @@ export abstract class ToggleableAction extends Action {
 	}
 }
 
-
-export interface IActionStateData {
-	className?: string;
-	label?: string;
-	tooltip?: string;
-	hideIcon?: boolean;
-	commandId?: string;
-}
-
-export class IMultiStateData<T> {
-	private _stateMap = new Map<T, IActionStateData>();
-	constructor(mappings: { key: T, value: IActionStateData }[], private _state: T, private _baseClass?: string) {
-		if (mappings) {
-			mappings.forEach(s => this._stateMap.set(s.key, s.value));
-		}
-	}
-
-	public set state(value: T) {
-		if (!this._stateMap.has(value)) {
-			throw new Error('State value must be in stateMap');
-		}
-		this._state = value;
-	}
-
-	public updateStateData(state: T, updater: (data: IActionStateData) => void): void {
-		let data = this._stateMap.get(state);
-		if (data) {
-			updater(data);
-		}
-	}
-
-	public get classes(): string {
-		let classVal = this.getStateValueOrDefault<string>((data) => data.className, '');
-		let classes = this._baseClass ? `${this._baseClass} ` : '';
-		classes += classVal;
-		if (this.getStateValueOrDefault<boolean>((data) => data.hideIcon, false)) {
-			classes += HIDE_ICON_CLASS;
-		}
-		return classes;
-	}
-
-	public get label(): string {
-		return this.getStateValueOrDefault<string>((data) => data.label, '');
-	}
-
-	public get tooltip(): string {
-		return this.getStateValueOrDefault<string>((data) => data.tooltip, '');
-	}
-
-	public get commandId(): string {
-		return this.getStateValueOrDefault<string>((data) => data.commandId, '');
-	}
-
-	private getStateValueOrDefault<U>(getter: (data: IActionStateData) => U, defaultVal?: U): U {
-		let data = this._stateMap.get(this._state);
-		return data ? getter(data) : defaultVal;
-	}
-}
-
-
-export abstract class MultiStateAction<T> extends Action {
-
-	constructor(
-		id: string,
-		protected states: IMultiStateData<T>,
-		private _keybindingService: IKeybindingService,
-		private readonly logService: ILogService) {
-		super(id, '');
-		this.updateLabelAndIcon();
-	}
-
-	private updateLabelAndIcon() {
-		let keyboardShortcut: string;
-		try {
-			// If a keyboard shortcut exists for the command id passed in, append that to the label
-			if (this.states.commandId !== '') {
-				let binding = this._keybindingService.lookupKeybinding(this.states.commandId);
-				keyboardShortcut = binding ? binding.getLabel() : undefined;
-			}
-		} catch (error) {
-			this.logService.error(error);
-		}
-		this.label = this.states.label;
-		this.tooltip = keyboardShortcut ? this.states.tooltip + ` (${keyboardShortcut})` : this.states.tooltip;
-		this.class = this.states.classes;
-	}
-
-	protected updateState(state: T): void {
-		this.states.state = state;
-		this.updateLabelAndIcon();
-	}
-}
-
 export class TrustedAction extends ToggleableAction {
 	// Constants
 	private static readonly trustedLabel = localize('trustLabel', "Trusted");
@@ -231,7 +134,7 @@ export class TrustedAction extends ToggleableAction {
 		this.toggle(value);
 	}
 
-	public run(context: NotebookComponent): Promise<boolean> {
+	public run(context: INotebookEditor): Promise<boolean> {
 		let self = this;
 		return new Promise<boolean>((resolve, reject) => {
 			try {
@@ -259,7 +162,7 @@ export class RunAllCellsAction extends Action {
 	) {
 		super(id, label, cssClass);
 	}
-	public async run(context: NotebookComponent): Promise<boolean> {
+	public async run(context: INotebookEditor): Promise<boolean> {
 		try {
 			await context.runAllCells();
 			return true;
@@ -291,15 +194,15 @@ export class CollapseCellsAction extends ToggleableAction {
 	public get isCollapsed(): boolean {
 		return this.state.isOn;
 	}
-	public set isCollapsed(value: boolean) {
+	private setCollapsed(value: boolean) {
 		this.toggle(value);
 	}
 
-	public run(context: NotebookComponent): Promise<boolean> {
+	public run(context: INotebookEditor): Promise<boolean> {
 		let self = this;
 		return new Promise<boolean>((resolve, reject) => {
 			try {
-				self.isCollapsed = !self.isCollapsed;
+				self.setCollapsed(!self.isCollapsed);
 				context.cells.forEach(cell => {
 					cell.isCollapsed = self.isCollapsed;
 				});
@@ -421,7 +324,7 @@ export class AttachToDropdown extends SelectBox {
 	}
 
 	// Load "Attach To" dropdown with the values corresponding to Kernel dropdown
-	public async loadAttachToDropdown(model: INotebookModel, currentKernel: string, showSelectConnection?: boolean): Promise<void> {
+	public loadAttachToDropdown(model: INotebookModel, currentKernel: string, showSelectConnection?: boolean): void {
 		let connProviderIds = this.model.getApplicableConnectionProviderIds(currentKernel);
 		if ((connProviderIds && connProviderIds.length === 0) || currentKernel === noKernel) {
 			this.setOptions([msgLocalHost]);
@@ -438,9 +341,9 @@ export class AttachToDropdown extends SelectBox {
 
 	public doChangeContext(connection?: ConnectionProfile, hideErrorMessage?: boolean): void {
 		if (this.value === msgChangeConnection || this.value === msgSelectConnection) {
-			this.openConnectionDialog();
+			this.openConnectionDialog().catch(err => this._notificationService.error(getErrorMessage(err)));
 		} else {
-			this.model.changeContext(this.value, connection, hideErrorMessage).then(ok => undefined, err => this._notificationService.error(getErrorMessage(err)));
+			this.model.changeContext(this.value, connection, hideErrorMessage).catch(err => this._notificationService.error(getErrorMessage(err)));
 		}
 	}
 
@@ -506,7 +409,7 @@ export class NewNotebookAction extends Action {
 	public static readonly ID = 'notebook.command.new';
 	public static readonly LABEL = localize('newNotebookAction', "New Notebook");
 
-	private static readonly INTERNAL_NEW_NOTEBOOK_CMD_ID = '_notebook.command.new';
+	public static readonly INTERNAL_NEW_NOTEBOOK_CMD_ID = '_notebook.command.new';
 	constructor(
 		id: string,
 		label: string,
@@ -527,5 +430,4 @@ export class NewNotebookAction extends Action {
 		}
 		return this.commandService.executeCommand(NewNotebookAction.INTERNAL_NEW_NOTEBOOK_CMD_ID, { connectionProfile: connProfile });
 	}
-
 }

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

@@ -0,0 +1,60 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { INotebookEditor, INotebookSection, INotebookParams } from 'sql/workbench/services/notebook/browser/notebookService';
+import { ICellModel, INotebookModel } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
+import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { CellType } from 'sql/workbench/contrib/notebook/common/models/contracts';
+
+export class NotebookComponentStub implements INotebookEditor {
+	get notebookParams(): INotebookParams {
+		throw new Error('Method not implemented.');
+	}
+	get id(): string {
+		throw new Error('Method not implemented.');
+	}
+	get cells(): ICellModel[] {
+		throw new Error('Method not implemented.');
+	}
+	get modelReady(): Promise<INotebookModel> {
+		throw new Error('Method not implemented.');
+	}
+	get model(): INotebookModel {
+		throw new Error('Method not implemented.');
+	}
+	isDirty(): boolean {
+		throw new Error('Method not implemented.');
+	}
+	isActive(): boolean {
+		throw new Error('Method not implemented.');
+	}
+	isVisible(): boolean {
+		throw new Error('Method not implemented.');
+	}
+	executeEdits(edits: ISingleNotebookEditOperation[]): boolean {
+		throw new Error('Method not implemented.');
+	}
+	runCell(cell: ICellModel): Promise<boolean> {
+		throw new Error('Method not implemented.');
+	}
+	runAllCells(startCell?: ICellModel, endCell?: ICellModel): Promise<boolean> {
+		throw new Error('Method not implemented.');
+	}
+	clearOutput(cell: ICellModel): Promise<boolean> {
+		throw new Error('Method not implemented.');
+	}
+	clearAllOutputs(): Promise<boolean> {
+		throw new Error('Method not implemented.');
+	}
+	getSections(): INotebookSection[] {
+		throw new Error('Method not implemented.');
+	}
+	navigateToSection(sectionId: string): void {
+		throw new Error('Method not implemented.');
+	}
+	addCell(cellType: CellType, index?: number, event?: Event) {
+		throw new Error('Method not implemented.');
+	}
+}

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

@@ -0,0 +1,156 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as TypeMoq from 'typemoq';
+import * as assert from 'assert';
+
+import { AddCellAction, ClearAllOutputsAction, CollapseCellsAction, TrustedAction, RunAllCellsAction, NewNotebookAction } from 'sql/workbench/contrib/notebook/browser/notebookActions';
+import { CellType } from 'sql/workbench/contrib/notebook/common/models/contracts';
+import { INotebookEditor } from 'sql/workbench/services/notebook/browser/notebookService';
+import { NotebookComponentStub } from 'sql/workbench/contrib/notebook/test/browser/common';
+import { ICellModel, INotebookModel } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
+import { INotificationService } from 'vs/platform/notification/common/notification';
+import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
+import { ICommandService } from 'vs/platform/commands/common/commands';
+import { TestCommandService } from 'vs/editor/test/browser/editorTestServices';
+
+suite('Notebook Actions', function (): void {
+	test('Add Cell Action', async function (): Promise<void> {
+		let testCellType: CellType = 'code';
+		let actualCellType: CellType;
+
+		let action = new AddCellAction('TestId', 'TestLabel', 'TestClass');
+		action.cellType = testCellType;
+
+		// Normal use case
+		let mockNotebookComponent = TypeMoq.Mock.ofType<INotebookEditor>(NotebookComponentStub);
+		mockNotebookComponent.setup(c => c.addCell(TypeMoq.It.isAny())).returns(cellType => {
+			actualCellType = cellType;
+		});
+
+		let result = await action.run(mockNotebookComponent.object);
+		assert.ok(result, 'Add Cell Action should succeed');
+		assert.strictEqual(actualCellType, testCellType);
+
+		// Handle error case
+		mockNotebookComponent.reset();
+		mockNotebookComponent.setup(c => c.addCell(TypeMoq.It.isAny())).throws(new Error('Test Error'));
+		await assert.rejects(action.run(mockNotebookComponent.object));
+	});
+
+	test('Clear All Outputs Action', async function (): Promise<void> {
+		let action = new ClearAllOutputsAction('TestId', 'TestLabel', 'TestClass');
+
+		// Normal use case
+		let mockNotebookComponent = TypeMoq.Mock.ofType<INotebookEditor>(NotebookComponentStub);
+		mockNotebookComponent.setup(c => c.clearAllOutputs()).returns(() => Promise.resolve(true));
+
+		let result = await action.run(mockNotebookComponent.object);
+		assert.ok(result, 'Clear All Outputs Action should succeed');
+		mockNotebookComponent.verify(c => c.clearAllOutputs(), TypeMoq.Times.once());
+
+		// Handle failure case
+		mockNotebookComponent.reset();
+		mockNotebookComponent.setup(c => c.clearAllOutputs()).returns(() => Promise.resolve(false));
+
+		result = await action.run(mockNotebookComponent.object);
+		assert.strictEqual(result, false, 'Clear All Outputs Action should have failed');
+		mockNotebookComponent.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', mockNotification.object);
+		assert.strictEqual(action.trusted, false, 'Should not be trusted by default');
+
+		// Normal use case
+		let contextStub = <INotebookEditor>{
+			model: <INotebookModel>{
+				trustedMode: false
+			}
+		};
+		let result = await action.run(contextStub);
+		assert.ok(result, 'Trusted Action should succeed');
+		assert.strictEqual(action.trusted, true, 'Should be trusted after toggling trusted state');
+
+		// Should stay trusted when trying to toggle again
+		result = await action.run(contextStub);
+		assert.ok(result, 'Trusted Action should succeed again');
+		assert.strictEqual(action.trusted, true, 'Should stay trusted when trying to toggle trusted to false');
+	});
+
+	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);
+
+		// Normal use case
+		let mockNotebookComponent = TypeMoq.Mock.ofType<INotebookEditor>(NotebookComponentStub);
+		mockNotebookComponent.setup(c => c.runAllCells()).returns(() => Promise.resolve(true));
+
+		let result = await action.run(mockNotebookComponent.object);
+		assert.ok(result, 'Run All Cells Action should succeed');
+		mockNotebookComponent.verify(c => c.runAllCells(), TypeMoq.Times.once());
+
+		// Handle errors
+		mockNotebookComponent.reset();
+		mockNotebookComponent.setup(c => c.runAllCells()).returns(() => { throw new Error('Test Error'); });
+
+		result = await action.run(mockNotebookComponent.object);
+		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');
+		assert.strictEqual(action.isCollapsed, false, 'Should not be collapsed by default');
+
+		let context = <INotebookEditor>{
+			cells: [<ICellModel>{
+				isCollapsed: false
+			}, <ICellModel>{
+				isCollapsed: true
+			}, <ICellModel>{
+				isCollapsed: false
+			}]
+		};
+
+		// Collapse cells case
+		let result = await action.run(context);
+		assert.ok(result, 'Collapse Cells Action should succeed');
+
+		assert.strictEqual(action.isCollapsed, true, 'Action should be collapsed after first toggle');
+		context.cells.forEach(cell => {
+			assert.strictEqual(cell.isCollapsed, true, 'Cells should be collapsed after first toggle');
+		});
+
+		// Toggle cells to uncollapsed
+		result = await action.run(context);
+		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 => {
+			assert.strictEqual(cell.isCollapsed, false, 'Cells should not be collapsed after second toggle');
+		});
+	});
+
+	test('New Notebook Action', async function (): Promise<void> {
+		let actualCmdId: string;
+
+		let mockCommandService = TypeMoq.Mock.ofType<ICommandService>(TestCommandService);
+		mockCommandService.setup(s => s.executeCommand(TypeMoq.It.isAny(), TypeMoq.It.isAny()))
+			.returns((commandId) => {
+				actualCmdId = commandId;
+				return Promise.resolve(true);
+			});
+
+		let action = new NewNotebookAction('TestId', 'TestLabel', mockCommandService.object, undefined);
+		action.run(undefined);
+
+		assert.strictEqual(actualCmdId, NewNotebookAction.INTERNAL_NEW_NOTEBOOK_CMD_ID);
+	});
+});

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

@@ -113,7 +113,8 @@ suite('Notebook Editor Model', function (): void {
 			clearOutput: undefined,
 			executeEdits: undefined,
 			getSections: undefined,
-			navigateToSection: undefined
+			navigateToSection: undefined,
+			addCell: undefined
 		};
 	});
 

src/sql/workbench/services/notebook/browser/notebookService.ts

@@ -5,7 +5,7 @@
 
 import * as azdata from 'azdata';
 
-import { Event } from 'vs/base/common/event';
+import * as vsEvent from 'vs/base/common/event';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { URI } from 'vs/base/common/uri';
 import { RenderMimeRegistry } from 'sql/workbench/contrib/notebook/browser/outputs/registry';
@@ -14,7 +14,7 @@ import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { NotebookInput } from 'sql/workbench/contrib/notebook/browser/models/notebookInput';
 import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { ICellModel, INotebookModel } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
-import { NotebookChangeType } from 'sql/workbench/contrib/notebook/common/models/contracts';
+import { NotebookChangeType, CellType } from 'sql/workbench/contrib/notebook/common/models/contracts';
 import { IBootstrapParams } from 'sql/workbench/services/bootstrap/common/bootstrapParams';
 
 export const SERVICE_ID = 'notebookService';
@@ -35,9 +35,9 @@ export interface ILanguageMagic {
 export interface INotebookService {
 	_serviceBrand: undefined;
 
-	readonly onNotebookEditorAdd: Event<INotebookEditor>;
-	readonly onNotebookEditorRemove: Event<INotebookEditor>;
-	onNotebookEditorRename: Event<INotebookEditor>;
+	readonly onNotebookEditorAdd: vsEvent.Event<INotebookEditor>;
+	readonly onNotebookEditorRemove: vsEvent.Event<INotebookEditor>;
+	onNotebookEditorRename: vsEvent.Event<INotebookEditor>;
 
 	readonly isRegistrationComplete: boolean;
 	readonly registrationComplete: Promise<void>;
@@ -159,6 +159,7 @@ export interface INotebookEditor {
 	clearAllOutputs(): Promise<boolean>;
 	getSections(): INotebookSection[];
 	navigateToSection(sectionId: string): void;
+	addCell(cellType: CellType, index?: number, event?: Event);
 }
 
 export interface INavigationProvider {