Add more notebook extension tests (#11143)

* more tests

* More tests

* More tests

* Add prompt tests

extensions/notebook/coverConfig.json

@@ -6,7 +6,9 @@
 		"**/integrationTest/**",
 		"**/node_modules/**",
 		"**/test/**",
-		"extension.js"
+		"extension.js",
+		"common/apiWrapper.js",
+		"common/localizedConstants.js"
 	],
 	"reports": [
 		"cobertura",

extensions/notebook/src/book/bookTreeView.ts

@@ -8,7 +8,7 @@ import * as vscode from 'vscode';
 import * as path from 'path';
 import * as fs from 'fs-extra';
 import * as constants from '../common/constants';
-import { IPrompter, QuestionTypes, IQuestion } from '../prompts/question';
+import { IPrompter, IQuestion, confirm } from '../prompts/question';
 import CodeAdapter from '../prompts/adapter';
 import { BookTreeItem, BookTreeItemType } from './bookTreeItem';
 import { BookModel } from './bookModel';
@@ -236,7 +236,7 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 		let notebookPath: string;
 		// If no uri is passed in, try to use the current active notebook editor
 		if (!uri) {
-			let openDocument = azdata.nb.activeNotebookEditor;
+			let openDocument = this._apiWrapper.getActiveNotebookEditor();
 			if (openDocument) {
 				notebookPath = openDocument.document.uri.fsPath;
 			}
@@ -511,10 +511,10 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 	//Confirmation message dialog
 	private async confirmReplace(): Promise<boolean> {
 		return await this.prompter.promptSingle<boolean>(<IQuestion>{
-			type: QuestionTypes.confirm,
+			type: confirm,
 			message: loc.confirmReplace,
 			default: false
-		});
+		}, this._apiWrapper);
 	}
 
 	getNavigation(uri: vscode.Uri): Thenable<azdata.nb.NavigationResult> {

extensions/notebook/src/common/apiWrapper.ts

@@ -69,6 +69,14 @@ export class ApiWrapper {
 		return azdata.nb.notebookDocuments;
 	}
 
+	public getActiveNotebookEditor(): azdata.nb.NotebookEditor {
+		return azdata.nb.activeNotebookEditor;
+	}
+
+	public showNotebookDocument(uri: vscode.Uri, showOptions?: azdata.nb.NotebookShowOptions): Thenable<azdata.nb.NotebookEditor> {
+		return azdata.nb.showNotebookDocument(uri, showOptions);
+	}
+
 	/**
 	 * Get the configuration for a extensionName
 	 * @param extensionName The string name of the extension to get the configuration for
@@ -95,4 +103,8 @@ export class ApiWrapper {
 	public createTreeView<T>(viewId: string, options: vscode.TreeViewOptions<T>): vscode.TreeView<T> {
 		return vscode.window.createTreeView(viewId, options);
 	}
+
+	public showQuickPick(items: string[] | Thenable<string[]>, options?: vscode.QuickPickOptions, token?: vscode.CancellationToken): Thenable<string | undefined> {
+		return vscode.window.showQuickPick(items, options, token);
+	}
 }

extensions/notebook/src/common/appContext.ts

@@ -12,19 +12,10 @@ import { NotebookUtils } from './notebookUtils';
  */
 export class AppContext {
 
-	private serviceMap: Map<string, any> = new Map();
 	public readonly notebookUtils: NotebookUtils;
 
 	constructor(public readonly extensionContext: vscode.ExtensionContext, public readonly apiWrapper: ApiWrapper) {
 		this.apiWrapper = apiWrapper || new ApiWrapper();
 		this.notebookUtils = new NotebookUtils(apiWrapper);
 	}
-
-	public getService<T>(serviceName: string): T {
-		return this.serviceMap.get(serviceName) as T;
-	}
-
-	public registerService<T>(serviceName: string, service: T): void {
-		this.serviceMap.set(serviceName, service);
-	}
 }

extensions/notebook/src/common/notebookUtils.ts

@@ -64,46 +64,46 @@ export class NotebookUtils {
 
 	public async runActiveCell(): Promise<void> {
 		try {
-			let notebook = azdata.nb.activeNotebookEditor;
+			let notebook = this._apiWrapper.getActiveNotebookEditor();
 			if (notebook) {
 				await notebook.runCell();
 			} else {
 				throw new Error(noNotebookVisible);
 			}
 		} catch (err) {
-			vscode.window.showErrorMessage(getErrorMessage(err));
+			this._apiWrapper.showErrorMessage(getErrorMessage(err));
 		}
 	}
 
 	public async clearActiveCellOutput(): Promise<void> {
 		try {
-			let notebook = azdata.nb.activeNotebookEditor;
+			let notebook = this._apiWrapper.getActiveNotebookEditor();
 			if (notebook) {
 				await notebook.clearOutput();
 			} else {
 				throw new Error(noNotebookVisible);
 			}
 		} catch (err) {
-			vscode.window.showErrorMessage(getErrorMessage(err));
+			this._apiWrapper.showErrorMessage(getErrorMessage(err));
 		}
 	}
 
 	public async runAllCells(startCell?: azdata.nb.NotebookCell, endCell?: azdata.nb.NotebookCell): Promise<void> {
 		try {
-			let notebook = azdata.nb.activeNotebookEditor;
+			let notebook = this._apiWrapper.getActiveNotebookEditor();
 			if (notebook) {
 				await notebook.runAllCells(startCell, endCell);
 			} else {
 				throw new Error(noNotebookVisible);
 			}
 		} catch (err) {
-			vscode.window.showErrorMessage(getErrorMessage(err));
+			this._apiWrapper.showErrorMessage(getErrorMessage(err));
 		}
 	}
 
 	public async addCell(cellType: azdata.nb.CellType): Promise<void> {
 		try {
-			let notebook = azdata.nb.activeNotebookEditor;
+			let notebook = this._apiWrapper.getActiveNotebookEditor();
 			if (notebook) {
 				await notebook.edit((editBuilder: azdata.nb.NotebookEditorEdit) => {
 					// TODO should prompt and handle cell placement
@@ -116,7 +116,7 @@ export class NotebookUtils {
 				throw new Error(noNotebookVisible);
 			}
 		} catch (err) {
-			vscode.window.showErrorMessage(getErrorMessage(err));
+			this._apiWrapper.showErrorMessage(getErrorMessage(err));
 		}
 	}
 
@@ -126,7 +126,7 @@ export class NotebookUtils {
 		let title = this.findNextUntitledEditorName();
 		let untitledUri = vscode.Uri.parse(`untitled:${title}`);
 
-		let editor = await azdata.nb.showNotebookDocument(untitledUri, {
+		let editor = await this._apiWrapper.showNotebookDocument(untitledUri, {
 			connectionProfile: oeContext ? oeContext.connectionProfile : undefined,
 			providerId: JUPYTER_NOTEBOOK_PROVIDER,
 			preview: false,
@@ -142,8 +142,7 @@ export class NotebookUtils {
 			if (hdfsPath.length > 0) {
 				let analyzeCommand = '#' + msgSampleCodeDataFrame + os.EOL + 'df = (spark.read.option("inferSchema", "true")'
 					+ os.EOL + '.option("header", "true")' + os.EOL + '.csv("{0}"))' + os.EOL + 'df.show(10)';
-
-				editor.edit(editBuilder => {
+				await editor.edit(editBuilder => {
 					editBuilder.insertCell({
 						cell_type: 'code',
 						source: analyzeCommand.replace('{0}', hdfsPath)

extensions/notebook/src/dialog/managePackages/installedPackagesTab.ts

@@ -10,7 +10,7 @@ import { JupyterServerInstallation, PythonPkgDetails } from '../../jupyter/jupyt
 import * as utils from '../../common/utils';
 import { ManagePackagesDialog } from './managePackagesDialog';
 import CodeAdapter from '../../prompts/adapter';
-import { QuestionTypes, IQuestion } from '../../prompts/question';
+import { IQuestion, confirm } from '../../prompts/question';
 
 const localize = nls.loadMessageBundle();
 
@@ -222,10 +222,10 @@ export class InstalledPackagesTab {
 
 		this.uninstallPackageButton.updateProperties({ enabled: false });
 		let doUninstall = await this.prompter.promptSingle<boolean>(<IQuestion>{
-			type: QuestionTypes.confirm,
+			type: confirm,
 			message: localize('managePackages.confirmUninstall', "Are you sure you want to uninstall the specified packages?"),
 			default: false
-		});
+		}, this.jupyterInstallation?.apiWrapper);
 
 		if (doUninstall) {
 			try {

extensions/notebook/src/jupyter/jupyterController.ts

@@ -15,7 +15,7 @@ import * as localizedConstants from '../common/localizedConstants';
 import { JupyterServerInstallation } from './jupyterServerInstallation';
 import { IServerInstance } from './common';
 import * as utils from '../common/utils';
-import { IPrompter, QuestionTypes, IQuestion } from '../prompts/question';
+import { IPrompter, IQuestion, confirm } from '../prompts/question';
 
 import { AppContext } from '../common/appContext';
 import { ApiWrapper } from '../common/apiWrapper';
@@ -42,6 +42,7 @@ export class JupyterController implements vscode.Disposable {
 
 	private outputChannel: vscode.OutputChannel;
 	private prompter: IPrompter;
+	private _notebookProvider: JupyterNotebookProvider;
 
 	constructor(private appContext: AppContext) {
 		this.prompter = new CodeAdapter();
@@ -56,6 +57,10 @@ export class JupyterController implements vscode.Disposable {
 		return this.appContext && this.appContext.extensionContext;
 	}
 
+	public get notebookProvider(): JupyterNotebookProvider {
+		return this._notebookProvider;
+	}
+
 	public dispose(): void {
 		this.deactivate();
 	}
@@ -89,23 +94,22 @@ export class JupyterController implements vscode.Disposable {
 		let supportedFileFilter: vscode.DocumentFilter[] = [
 			{ scheme: 'untitled', language: '*' }
 		];
-		let notebookProvider = this.registerNotebookProvider();
-		this.extensionContext.subscriptions.push(this.apiWrapper.registerCompletionItemProvider(supportedFileFilter, new NotebookCompletionItemProvider(notebookProvider)));
+		this.registerNotebookProvider();
+		this.extensionContext.subscriptions.push(this.apiWrapper.registerCompletionItemProvider(supportedFileFilter, new NotebookCompletionItemProvider(this._notebookProvider)));
 
 		this.registerDefaultPackageManageProviders();
 		return true;
 	}
 
-	private registerNotebookProvider(): JupyterNotebookProvider {
-		let notebookProvider = new JupyterNotebookProvider((documentUri: vscode.Uri) => new LocalJupyterServerManager({
+	private registerNotebookProvider(): void {
+		this._notebookProvider = new JupyterNotebookProvider((documentUri: vscode.Uri) => new LocalJupyterServerManager({
 			documentPath: documentUri.fsPath,
 			jupyterInstallation: this._jupyterInstallation,
 			extensionContext: this.extensionContext,
 			apiWrapper: this.apiWrapper,
 			factory: this._serverInstanceFactory
 		}));
-		azdata.nb.registerNotebookProvider(notebookProvider);
-		return notebookProvider;
+		azdata.nb.registerNotebookProvider(this._notebookProvider);
 	}
 
 	private saveProfileAndCreateNotebook(profile: azdata.IConnectionProfile): Promise<void> {
@@ -198,10 +202,10 @@ export class JupyterController implements vscode.Disposable {
 	//Confirmation message dialog
 	private async confirmReinstall(): Promise<boolean> {
 		return await this.prompter.promptSingle<boolean>(<IQuestion>{
-			type: QuestionTypes.confirm,
+			type: confirm,
 			message: localize('confirmReinstall', "Are you sure you want to reinstall?"),
 			default: true
-		});
+		}, this.apiWrapper);
 	}
 
 	public async doManagePackages(options?: ManagePackageDialogOptions): Promise<void> {

extensions/notebook/src/jupyter/jupyterNotebookProvider.ts

@@ -31,8 +31,12 @@ export class JupyterNotebookProvider implements nb.NotebookProvider {
 		return Promise.resolve(this.doGetNotebookManager(notebookUri));
 	}
 
+	public get notebookManagerCount(): number {
+		return this.managerTracker.size;
+	}
+
 	private doGetNotebookManager(notebookUri: vscode.Uri): nb.NotebookManager {
-		let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString());
+		let baseFolder = this.transformToBaseFolder(notebookUri?.fsPath?.toString());
 		let manager = this.managerTracker.get(baseFolder);
 		if (!manager) {
 			let baseFolderUri = vscode.Uri.file(baseFolder);

extensions/notebook/src/jupyter/jupyterServerInstallation.ts

@@ -18,7 +18,7 @@ import * as utils from '../common/utils';
 import { OutputChannel, ConfigurationTarget, window } from 'vscode';
 import { Deferred } from '../common/promise';
 import { ConfigurePythonWizard } from '../dialog/configurePython/configurePythonWizard';
-import { IPrompter, IQuestion, QuestionTypes } from '../prompts/question';
+import { IPrompter, IQuestion, confirm } from '../prompts/question';
 import CodeAdapter from '../prompts/adapter';
 import { ConfigurePythonDialog } from '../dialog/configurePython/configurePythonDialog';
 
@@ -587,10 +587,10 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 			let doUpgrade: boolean;
 			if (promptForUpgrade) {
 				doUpgrade = await this._prompter.promptSingle<boolean>(<IQuestion>{
-					type: QuestionTypes.confirm,
+					type: confirm,
 					message: localize('confirmPackageUpgrade', "Some required python packages need to be installed. Would you like to install them now?"),
 					default: true
-				});
+				}, this.apiWrapper);
 				if (!doUpgrade) {
 					throw new Error(localize('configurePython.packageInstallDeclined', "Package installation was declined."));
 				}

extensions/notebook/src/prompts/adapter.ts

@@ -1,32 +1,17 @@
 // This code is originally from https://github.com/DonJayamanne/bowerVSCode
 // License: https://github.com/DonJayamanne/bowerVSCode/blob/master/LICENSE
 
-import { window } from 'vscode';
 import PromptFactory from './factory';
 import EscapeException from './escapeException';
-import { IQuestion, IPrompter, IPromptCallback } from './question';
+import { IQuestion, IPrompter } from './question';
+import { ApiWrapper } from '../common/apiWrapper';
 
 // Supports simple pattern for prompting for user input and acting on this
 export default class CodeAdapter implements IPrompter {
 
-	// TODO define question interface
-	private fixQuestion(question: any): any {
-		if (question.type === 'checkbox' && Array.isArray(question.choices)) {
-			// For some reason when there's a choice of checkboxes, they aren't formatted properly
-			// Not sure where the issue is
-			question.choices = question.choices.map((item: any) => {
-				if (typeof (item) === 'string') {
-					return { checked: false, name: item, value: item };
-				} else {
-					return item;
-				}
-			});
-		}
-	}
-
-	public promptSingle<T>(question: IQuestion, ignoreFocusOut?: boolean): Promise<T> {
+	public promptSingle<T>(question: IQuestion, apiWrapper?: ApiWrapper): Promise<T> {
 		let questions: IQuestion[] = [question];
-		return this.prompt(questions, ignoreFocusOut).then((answers: { [key: string]: T }) => {
+		return this.prompt(questions, apiWrapper).then((answers: { [key: string]: T }) => {
 			if (answers) {
 				let response: T = answers[question.name];
 				return response || undefined;
@@ -35,15 +20,13 @@ export default class CodeAdapter implements IPrompter {
 		});
 	}
 
-	public prompt<T>(questions: IQuestion[], ignoreFocusOut?: boolean): Promise<{ [key: string]: T }> {
+	public prompt<T>(questions: IQuestion[], apiWrapper = new ApiWrapper()): Promise<{ [key: string]: T }> {
 		let answers: { [key: string]: T } = {};
 
 		// Collapse multiple questions into a set of prompt steps
 		let promptResult: Promise<{ [key: string]: T }> = questions.reduce((promise: Promise<{ [key: string]: T }>, question: IQuestion) => {
-			this.fixQuestion(question);
-
 			return promise.then(() => {
-				return PromptFactory.createPrompt(question, ignoreFocusOut);
+				return PromptFactory.createPrompt(question);
 			}).then(prompt => {
 				// Original Code: uses jQuery patterns. Keeping for reference
 				// if (!question.when || question.when(answers) === true) {
@@ -53,7 +36,7 @@ export default class CodeAdapter implements IPrompter {
 				// }
 
 				if (!question.shouldPrompt || question.shouldPrompt(answers) === true) {
-					return prompt.render().then((result: any) => {
+					return prompt.render(apiWrapper).then((result: any) => {
 						answers[question.name] = result;
 
 						if (question.onAnswered) {
@@ -71,17 +54,7 @@ export default class CodeAdapter implements IPrompter {
 				return undefined;
 			}
 
-			window.showErrorMessage(err.message);
-		});
-	}
-
-	// Helper to make it possible to prompt using callback pattern. Generally Promise is a preferred flow
-	public promptCallback(questions: IQuestion[], callback: IPromptCallback): void {
-		// Collapse multiple questions into a set of prompt steps
-		this.prompt(questions).then(answers => {
-			if (callback) {
-				callback(answers);
-			}
+			apiWrapper.showErrorMessage(err.message);
 		});
 	}
 }

extensions/notebook/src/prompts/confirm.ts

@@ -3,18 +3,18 @@
 // This code is originally from https://github.com/DonJayamanne/bowerVSCode
 // License: https://github.com/DonJayamanne/bowerVSCode/blob/master/LICENSE
 
-import { window } from 'vscode';
 import Prompt from './prompt';
 import LocalizedConstants = require('../common/localizedConstants');
 import EscapeException from './escapeException';
+import { ApiWrapper } from '../common/apiWrapper';
 
 export default class ConfirmPrompt extends Prompt {
 
-	constructor(question: any, ignoreFocusOut?: boolean) {
-		super(question, ignoreFocusOut);
+	constructor(question: any) {
+		super(question);
 	}
 
-	public render(): any {
+	public render(apiWrapper: ApiWrapper): any {
 		let choices: { [id: string]: boolean } = {};
 		choices[LocalizedConstants.msgYes] = true;
 		choices[LocalizedConstants.msgNo] = false;
@@ -22,7 +22,7 @@ export default class ConfirmPrompt extends Prompt {
 		let options = this.defaultQuickPickOptions;
 		options.placeHolder = this._question.message;
 
-		return window.showQuickPick(Object.keys(choices), options)
+		return apiWrapper.showQuickPick(Object.keys(choices), options)
 			.then(result => {
 				if (result === undefined) {
 					throw new EscapeException();

extensions/notebook/src/prompts/factory.ts

@@ -8,10 +8,10 @@ import ConfirmPrompt from './confirm';
 
 export default class PromptFactory {
 
-	public static createPrompt(question: any, ignoreFocusOut?: boolean): Prompt {
+	public static createPrompt(question: any): Prompt {
 		switch (question.type) {
 			case 'confirm':
-				return new ConfirmPrompt(question, ignoreFocusOut);
+				return new ConfirmPrompt(question);
 			default:
 				throw new Error(`Could not find a prompt for question type ${question.type}`);
 		}

extensions/notebook/src/prompts/prompt.ts

@@ -4,6 +4,7 @@
 // License: https://github.com/DonJayamanne/bowerVSCode/blob/master/LICENSE
 
 import { InputBoxOptions, QuickPickOptions } from 'vscode';
+import { ApiWrapper } from '../common/apiWrapper';
 
 abstract class Prompt {
 
@@ -15,7 +16,7 @@ abstract class Prompt {
 		this._ignoreFocusOut = ignoreFocusOut ? ignoreFocusOut : false;
 	}
 
-	public abstract render(): any;
+	public abstract render(apiWrapper: ApiWrapper): any;
 
 	protected get defaultQuickPickOptions(): QuickPickOptions {
 		return {

extensions/notebook/src/prompts/question.ts

@@ -4,15 +4,9 @@
  *--------------------------------------------------------------------------------------------*/
 
 import vscode = require('vscode');
+import { ApiWrapper } from '../common/apiWrapper';
 
-export class QuestionTypes {
-	public static get input(): string { return 'input'; }
-	public static get password(): string { return 'password'; }
-	public static get list(): string { return 'list'; }
-	public static get confirm(): string { return 'confirm'; }
-	public static get checkbox(): string { return 'checkbox'; }
-	public static get expand(): string { return 'expand'; }
-}
+export const confirm = 'confirm';
 
 // Question interface to clarify how to use the prompt feature
 // based on Bower Question format: https://github.com/bower/bower/blob/89069784bb46bfd6639b4a75e98a0d7399a8c2cb/packages/bower-logger/README.md
@@ -54,15 +48,14 @@ export interface IQuestionHandler {
 }
 
 export interface IPrompter {
-	promptSingle<T>(question: IQuestion, ignoreFocusOut?: boolean): Promise<T>;
+	promptSingle<T>(question: IQuestion, apiWrapper?: ApiWrapper): Promise<T>;
 	/**
 	 * Prompts for multiple questions
 	 *
 	 * @returns Map of question IDs to results, or undefined if
 	 * the user canceled the question session
 	 */
-	prompt<T>(questions: IQuestion[], ignoreFocusOut?: boolean): Promise<{ [questionId: string]: any }>;
-	promptCallback(questions: IQuestion[], callback: IPromptCallback): void;
+	prompt<T>(questions: IQuestion[], apiWrapper?: ApiWrapper): Promise<{ [questionId: string]: any }>;
 }
 
 export interface IPromptCallback {

extensions/notebook/src/protocol/notebookUriHandler.ts

@@ -11,7 +11,7 @@ import * as path from 'path';
 import * as querystring from 'querystring';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
-import { IQuestion, QuestionTypes } from '../prompts/question';
+import { IQuestion, confirm } from '../prompts/question';
 import CodeAdapter from '../prompts/adapter';
 import { getErrorMessage, isEditorTitleFree } from '../common/utils';
 
@@ -61,7 +61,7 @@ export class NotebookUriHandler implements vscode.UriHandler {
 			}
 
 			let doOpen = await this.prompter.promptSingle<boolean>(<IQuestion>{
-				type: QuestionTypes.confirm,
+				type: confirm,
 				message: localize('notebook.confirmOpen', "Download and open '{0}'?", url),
 				default: true
 			});

extensions/notebook/src/test/common/notebookUtils.test.ts

@@ -14,15 +14,21 @@ import * as uuid from 'uuid';
 import { promises as fs } from 'fs';
 import { ApiWrapper } from '../../common/apiWrapper';
 import { tryDeleteFile } from './testUtils';
+import { CellTypes } from '../../contracts/content';
 
 describe('notebookUtils Tests', function (): void {
 	let notebookUtils: NotebookUtils;
 	let apiWrapperMock: TypeMoq.IMock<ApiWrapper>;
-	before(function (): void {
+
+	beforeEach(function (): void {
 		apiWrapperMock = TypeMoq.Mock.ofInstance(new ApiWrapper());
 		notebookUtils = new NotebookUtils(apiWrapperMock.object);
 	});
 
+	this.afterAll(async function (): Promise<void> {
+		await vscode.commands.executeCommand('workbench.action.closeAllEditors');
+	});
+
 	describe('newNotebook', function (): void {
 		it('Should open a new notebook successfully', async function (): Promise<void> {
 			should(azdata.nb.notebookDocuments.length).equal(0, 'There should be not any open Notebook documents');
@@ -80,4 +86,117 @@ describe('notebookUtils Tests', function (): void {
 			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
 		});
 	});
+
+	describe('runActiveCell', function () {
+		it('shows error when no notebook visible', async function (): Promise<void> {
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => undefined);
+			await notebookUtils.runActiveCell();
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+
+		it('does not show error when notebook visible', async function (): Promise<void> {
+			let mockNotebookEditor = TypeMoq.Mock.ofType<azdata.nb.NotebookEditor>();
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => mockNotebookEditor.object);
+			await notebookUtils.runActiveCell();
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.never());
+			mockNotebookEditor.verify(x => x.runCell(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+	});
+
+	describe('clearActiveCellOutput', function () {
+		it('shows error when no notebook visible', async function (): Promise<void> {
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => undefined);
+			await notebookUtils.clearActiveCellOutput();
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+
+		it('does not show error when notebook visible', async function (): Promise<void> {
+			let mockNotebookEditor = TypeMoq.Mock.ofType<azdata.nb.NotebookEditor>();
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => mockNotebookEditor.object);
+			await notebookUtils.clearActiveCellOutput();
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.never());
+			mockNotebookEditor.verify(x => x.clearOutput(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+	});
+
+	describe('runAllCells', function () {
+		it('shows error when no notebook visible', async function (): Promise<void> {
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => undefined);
+			await notebookUtils.runAllCells();
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+
+		it('does not show error when notebook visible', async function (): Promise<void> {
+			let mockNotebookEditor = TypeMoq.Mock.ofType<azdata.nb.NotebookEditor>();
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => mockNotebookEditor.object);
+			await notebookUtils.runAllCells();
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.never());
+			mockNotebookEditor.verify(x => x.runAllCells(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+	});
+
+	describe('addCell', function () {
+		it('shows error when no notebook visible for code cell', async function (): Promise<void> {
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => undefined);
+			await notebookUtils.addCell('code');
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+
+		it('shows error when no notebook visible for markdown cell', async function (): Promise<void> {
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => undefined);
+			await notebookUtils.addCell('markdown');
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		});
+
+		it('does not show error when notebook visible for code cell', async function (): Promise<void> {
+			const notebookEditor = await notebookUtils.newNotebook(undefined);
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => notebookEditor);
+			await notebookUtils.addCell('code');
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.never());
+			should(notebookEditor.document.cells.length).equal(1);
+			should(notebookEditor.document.cells[0].contents.cell_type).equal(CellTypes.Code);
+		});
+
+		it('does not show error when notebook visible for markdown cell', async function (): Promise<void> {
+			const notebookEditor = await notebookUtils.newNotebook(undefined);
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => notebookEditor);
+			await notebookUtils.addCell('markdown');
+			apiWrapperMock.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.never());
+			should(notebookEditor.document.cells.length).equal(1);
+			should(notebookEditor.document.cells[0].contents.cell_type).equal(CellTypes.Markdown);
+		});
+	});
+
+	describe('analyzeNotebook', function () {
+		it('creates cell when oeContext exists', async function (): Promise<void> {
+			const notebookEditor = await notebookUtils.newNotebook(undefined);
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => notebookEditor);
+			apiWrapperMock.setup(x => x.showNotebookDocument(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(notebookEditor));
+			const oeContext: azdata.ObjectExplorerContext = {
+				connectionProfile: undefined,
+				isConnectionNode: true,
+				nodeInfo: {
+					nodePath: 'path/HDFS/path2',
+					errorMessage: undefined,
+					isLeaf: false,
+					label: 'fakeLabel',
+					metadata: undefined,
+					nodeStatus: undefined,
+					nodeSubType: undefined,
+					nodeType: undefined
+				}
+			}
+			await notebookUtils.analyzeNotebook(oeContext);
+			should(notebookEditor.document.cells.length).equal(1, 'One cell should exist');
+			should(notebookEditor.document.cells[0].contents.cell_type).equal(CellTypes.Code, 'Cell was created with incorrect type');
+		});
+
+		it('does not create new cell when oeContext does not exist', async function (): Promise<void> {
+			const notebookEditor = await notebookUtils.newNotebook(undefined);
+			apiWrapperMock.setup(x => x.getActiveNotebookEditor()).returns(() => notebookEditor);
+			apiWrapperMock.setup(x => x.showNotebookDocument(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(notebookEditor));
+			await notebookUtils.analyzeNotebook();
+			should(notebookEditor.document.cells.length).equal(0, 'No cells should exist');
+		});
+	});
 });

extensions/notebook/src/test/common/prompt.test.ts

@@ -0,0 +1,50 @@
+/*---------------------------------------------------------------------------------------------
+ *  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 { IPrompter, confirm, IQuestion } from '../../prompts/question';
+import CodeAdapter from '../../prompts/adapter';
+import { ApiWrapper } from '../../common/apiWrapper';
+
+
+describe('Prompt', () => {
+
+	let prompter: IPrompter;
+	let mockApiWrapper: TypeMoq.IMock<ApiWrapper>;
+
+
+	before(function () {
+		prompter = new CodeAdapter();
+	});
+
+	beforeEach(function () {
+		mockApiWrapper = TypeMoq.Mock.ofType<ApiWrapper>();
+		mockApiWrapper.setup(x => x.showErrorMessage(TypeMoq.It.isAny()));
+	});
+
+	it('Should not find prompt for invalid question type', async function (): Promise<void> {
+		await prompter.promptSingle<boolean>(<IQuestion>{
+			type: 'invalidType',
+			message: 'sample message',
+			default: false
+		}, mockApiWrapper.object);
+		mockApiWrapper.verify(s => s.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		mockApiWrapper.verify(s => s.showQuickPick(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.never());
+	});
+
+	it('Should find prompt for confirm type', async function (): Promise<void> {
+		mockApiWrapper.setup(x => x.showQuickPick(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve('Yes'));
+		await prompter.promptSingle<boolean>(<IQuestion>{
+			type: confirm,
+			message: 'sample message',
+			default: false
+		}, mockApiWrapper.object);
+		mockApiWrapper.verify(s => s.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.never());
+		mockApiWrapper.verify(s => s.showQuickPick(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
+	});
+});
+
+

extensions/notebook/src/test/model/jupyterController.test.ts

@@ -3,6 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import * as should from 'should';
 import * as TypeMoq from 'typemoq';
@@ -12,28 +13,34 @@ import { AppContext } from '../../common/appContext';
 import { JupyterController } from '../../jupyter/jupyterController';
 import { LocalPipPackageManageProvider } from '../../jupyter/localPipPackageManageProvider';
 import { MockExtensionContext } from '../common/stubs';
+import { NotebookUtils } from '../../common/notebookUtils';
 
-describe('JupyterController tests', function () {
+describe('Jupyter Controller', function () {
 	let mockExtensionContext: vscode.ExtensionContext;
 	let appContext: AppContext;
 	let controller: JupyterController;
 	let mockApiWrapper: TypeMoq.IMock<ApiWrapper>;
+	let connection: azdata.connection.ConnectionProfile;
 
 	this.beforeAll(() => {
 		mockExtensionContext = new MockExtensionContext();
-		mockApiWrapper = TypeMoq.Mock.ofType<ApiWrapper>();
-		appContext = new AppContext(mockExtensionContext, mockApiWrapper.object);
+		connection  = new azdata.connection.ConnectionProfile();
 	});
 
 	this.beforeEach(() => {
+		mockApiWrapper = TypeMoq.Mock.ofType<ApiWrapper>();
+		mockApiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
+		appContext = new AppContext(mockExtensionContext, mockApiWrapper.object);
 		controller = new JupyterController(appContext);
 	});
 
 	it('should activate new JupyterController successfully', async () => {
 		should(controller.extensionContext).deepEqual(appContext.extensionContext, 'Extension context should be passed through');
+		should(controller.jupyterInstallation).equal(undefined, 'JupyterInstallation should be undefined before controller activation');
 		await should(controller.activate()).not.be.rejected();
 		// On activation, local pip and local conda package providers should exist
 		should(controller.packageManageProviders.size).equal(2, 'Local pip and conda package providers should be default providers');
+		should(controller.jupyterInstallation.extensionPath).equal(appContext.extensionContext.extensionPath, 'JupyterInstallation extension path should match appContext extension path');
 	});
 
 	it('should create new packageManageProvider successfully', async () => {
@@ -50,4 +57,43 @@ describe('JupyterController tests', function () {
 		should.throws(() => controller.registerPackageManager('provider1', mockProvider.object));
 		should(controller.packageManageProviders.size).equal(1, 'Package manage providers should still equal 1');
 	});
+
+	it('should should get defaultConnection() successfully', async () => {
+		let defaultConnection = await controller.getDefaultConnection();
+		should(defaultConnection).deepEqual(connection, 'getDefaultConnection() did not return expected result');
+	});
+
+	it('should show error message for doManagePackages before activation', async () => {
+		await controller.doManagePackages();
+		mockApiWrapper.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.once());
+	});
+
+	it('should not show error message for doManagePackages after activation', async () => {
+		await controller.activate();
+		await controller.doManagePackages();
+		mockApiWrapper.verify(x => x.showErrorMessage(TypeMoq.It.isAny()), TypeMoq.Times.never());
+	});
+
+	it('Returns expected values from notebook provider', async () =>  {
+		await controller.activate();
+		should(controller.notebookProvider.standardKernels).deepEqual([], 'Notebook provider standard kernels should return empty array');
+		should(controller.notebookProvider.providerId).equal('jupyter', 'Notebook provider should be jupyter');
+		should(controller.notebookProvider.getNotebookManager(undefined)).be.rejected();
+		should(controller.notebookProvider.notebookManagerCount).equal(0);
+		controller.notebookProvider.handleNotebookClosed(undefined);
+	});
+
+	it('Returns notebook manager for real notebook editor', async () =>  {
+		await controller.activate();
+		let notebookUtils = new NotebookUtils(mockApiWrapper.object);
+		const notebookEditor = await notebookUtils.newNotebook(undefined);
+		let notebookManager = await controller.notebookProvider.getNotebookManager(notebookEditor.document.uri);
+		should(controller.notebookProvider.notebookManagerCount).equal(1);
+
+		// Session manager should not be immediately ready
+		should(notebookManager.sessionManager.isReady).equal(false);
+		// Session manager should not immediately have specs
+		should(notebookManager.sessionManager.specs).equal(undefined);
+		controller.notebookProvider.handleNotebookClosed(notebookEditor.document.uri);
+	});
 });