Fix #4356 New Notebook from connection doesn't connect (#4364)

* Fix #4356 New Notebook from connection doesn't connect
Fix new notebook error by passing profile instead of ID.
- I could've just sent the ID over, but this fix sets the stage for disconnected connections to work (since we have enough info to properly connect).
- There's a bug in NotebookModel blocking the disconnected connection part working, but Yurong's in progress fixes will unblock this. Hence checking in as-is and working to properly unblock once that's in.

* Support connection profile in commandline service
- Added new context API for things that want to work on commandline and object explorer
- Refactored commandlineservice slightly to be async & have a simpler execution flow (far fewer if/else statements)

* Fix unit tests
- Fixed 2 issues raised by tests (sholdn't do new query if no profile passed, shouldn't error on new query failing)
- Updated unit tests to pass as expected given changes to the APIs.

extensions/mssql/src/main.ts

@@ -177,7 +177,7 @@ async function handleNewNotebookTask(oeContext?: azdata.ObjectExplorerContext, p
 	// to handle this. We should look into improving this in the future
 	let untitledUri = vscode.Uri.parse(`untitled:Notebook-${untitledCounter++}`);
 	let editor = await azdata.nb.showNotebookDocument(untitledUri, {
-		connectionId: profile.id,
+		connectionProfile: profile,
 		providerId: jupyterNotebookProviderId,
 		preview: false,
 		defaultKernel: {
@@ -220,7 +220,7 @@ async function handleOpenNotebookTask(profile: azdata.IConnectionProfile): Promi
 			vscode.window.showErrorMessage(localize('unsupportedFileType', 'Only .ipynb Notebooks are supported'));
 		} else {
 			await azdata.nb.showNotebookDocument(fileUri, {
-				connectionId: profile.id,
+				connectionProfile: profile,
 				providerId: jupyterNotebookProviderId,
 				preview: false
 			});

extensions/notebook/src/extension.ts

@@ -25,8 +25,12 @@ let counter = 0;
 export let controller: JupyterController;
 
 export function activate(extensionContext: vscode.ExtensionContext) {
-	extensionContext.subscriptions.push(vscode.commands.registerCommand('notebook.command.new', (connectionId?: string) => {
-		newNotebook(connectionId);
+	extensionContext.subscriptions.push(vscode.commands.registerCommand('notebook.command.new', (context?: azdata.ConnectedContext) => {
+		let connectionProfile: azdata.IConnectionProfile = undefined;
+		if (context && context.connectionProfile) {
+			connectionProfile = context.connectionProfile;
+		}
+		newNotebook(connectionProfile);
 	}));
 	extensionContext.subscriptions.push(vscode.commands.registerCommand('notebook.command.open', () => {
 		openNotebook();
@@ -49,15 +53,15 @@ export function activate(extensionContext: vscode.ExtensionContext) {
 	controller.activate();
 }
 
-function newNotebook(connectionId: string) {
+function newNotebook(connectionProfile: azdata.IConnectionProfile) {
 	let title = `Untitled-${counter++}`;
 	let untitledUri = vscode.Uri.parse(`untitled:${title}`);
-	let options: azdata.nb.NotebookShowOptions = connectionId ? {
+	let options: azdata.nb.NotebookShowOptions = connectionProfile ? {
 		viewColumn: null,
 		preserveFocus: true,
 		preview: null,
 		providerId: null,
-		connectionId: connectionId,
+		connectionProfile: connectionProfile,
 		defaultKernel: null
 	} : null;
 	azdata.nb.showNotebookDocument(untitledUri, options).then(success => {
@@ -122,7 +126,7 @@ async function analyzeNotebook(oeContext?: azdata.ObjectExplorerContext): Promis
 	let untitledUri = vscode.Uri.parse(`untitled:Notebook-${counter++}`);
 
 	let editor = await azdata.nb.showNotebookDocument(untitledUri, {
-		connectionId: oeContext ? oeContext.connectionProfile.id : '',
+		connectionProfile: oeContext ? oeContext.connectionProfile : undefined,
 		providerId: JUPYTER_NOTEBOOK_PROVIDER,
 		preview: false,
 		defaultKernel: {

extensions/notebook/src/integrationTest/notebookIntegration.test.ts

@@ -65,11 +65,11 @@ describe('Notebook Integration Test', function (): void {
 		await ensureJupyterInstalled();
 
 		// Given a connection to a server exists
-		let connectionId = await connectToSparkIntegrationServer();
+		let connectionProfile = await connectToSparkIntegrationServer();
 
 		// When I open a Spark notebook and run the cell
 		let notebook = await azdata.nb.showNotebookDocument(uri, {
-			connectionId: connectionId
+			connectionProfile: connectionProfile
 		});
 		should(notebook.document.cells).have.length(1);
 		let ran = await notebook.runCell(notebook.document.cells[0]);
@@ -90,7 +90,7 @@ describe('Notebook Integration Test', function (): void {
 	});
 });
 
-async function connectToSparkIntegrationServer(): Promise<string> {
+async function connectToSparkIntegrationServer(): Promise<azdata.IConnectionProfile> {
 	assert.ok(process.env.BACKEND_HOSTNAME, 'BACKEND_HOSTNAME, BACKEND_USERNAME, BACKEND_PWD must be set using ./tasks/setbackenvariables.sh or .\\tasks\\setbackendvaraibles.bat');
 	let connInfo: azdata.connection.Connection = {
 		options: {
@@ -114,7 +114,7 @@ async function connectToSparkIntegrationServer(): Promise<string> {
 	let activeConnections = await azdata.connection.getActiveConnections();
 	should(activeConnections).have.length(1);
 
-	return result.connectionId;
+	return <azdata.IConnectionProfile><any>connInfo;
 }
 
 function writeNotebookToFile(pythonNotebook: INotebook): vscode.Uri {

extensions/notebook/src/jupyter/jupyterController.ts

@@ -152,7 +152,7 @@ export class JupyterController implements vscode.Disposable {
 				this.apiWrapper.showErrorMessage(localize('unsupportedFileType', 'Only .ipynb Notebooks are supported'));
 			} else {
 				await azdata.nb.showNotebookDocument(fileUri, {
-					connectionId: profile.id,
+					connectionProfile: profile,
 					providerId: constants.jupyterNotebookProviderId,
 					preview: false
 				});
@@ -165,7 +165,7 @@ export class JupyterController implements vscode.Disposable {
 		// to handle this. We should look into improving this in the future
 		let untitledUri = vscode.Uri.parse(`untitled:Notebook-${untitledCounter++}`);
 		let editor = await azdata.nb.showNotebookDocument(untitledUri, {
-			connectionId: profile.id,
+			connectionProfile: profile,
 			providerId: constants.jupyterNotebookProviderId,
 			preview: false,
 			defaultKernel: {

src/sql/azdata.proposed.d.ts

@@ -3667,19 +3667,30 @@ declare module 'azdata' {
 		export function getProvidersByType<T extends DataProvider>(providerType: DataProviderType): T[];
 	}
 
+
 	/**
 	 * Context object passed as an argument to command callbacks.
-	 * Defines the key properties required to identify a node in the object
-	 * explorer tree and take action against it.
+	 * Defines properties that can be sent for any connected context,
+	 * whether that is the Object Explorer context menu or a command line
+	 * startup argument.
 	 */
-	export interface ObjectExplorerContext {
 
+	export interface ConnectedContext {
 		/**
 		 * The connection information for the selected object.
 		 * Note that the connection is not guaranteed to be in a connected
 		 * state on click.
 		 */
 		connectionProfile: IConnectionProfile;
+	}
+
+	/**
+	 * Context object passed as an argument to command callbacks.
+	 * Defines the key properties required to identify a node in the object
+	 * explorer tree and take action against it.
+	 */
+	export interface ObjectExplorerContext extends ConnectedContext {
+
 		/**
 		 * Defines whether this is a Connection-level object.
 		 * If not, the object is expected to be a child object underneath
@@ -4016,9 +4027,9 @@ declare module 'azdata' {
 			providerId?: string;
 
 			/**
-			 * Optional ID indicating the initial connection to use for this editor
+			 * Optional profile indicating the initial connection to use for this editor
 			 */
-			connectionId?: string;
+			connectionProfile?: IConnectionProfile;
 
 			/**
 			 * Default kernel for notebook

src/sql/parts/common/customInputConverter.ts

@@ -59,7 +59,6 @@ export function convertEditorInput(input: EditorInput, options: IQueryEditorOpti
 		if (uri) {
 			return withService<INotebookService, NotebookInput>(instantiationService, INotebookService, notebookService => {
 				let fileName: string = 'untitled';
-				let providerIds: string[] = [DEFAULT_NOTEBOOK_PROVIDER];
 				if (input) {
 					fileName = input.getName();
 				}

src/sql/parts/notebook/models/notebookModel.ts

@@ -408,6 +408,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			}
 			let profile = new ConnectionProfile(this._notebookOptions.capabilitiesService, this.connectionProfile);
 
+			// TODO: this code needs to be fixed since it is called before the this._savedKernelInfo is set.
+			// This means it always fails, and we end up using the default connection instead. If you right-click
+			// and run "New Notebook" on a disconnected server this means you get the wrong connection (global active)
+			// instead of the one you chose, or it'll fail to connect in general
 			if (this.isValidConnection(profile)) {
 				this._activeConnection = profile;
 			} else {

src/sql/parts/notebook/notebook.component.ts

@@ -104,23 +104,18 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 	private updateProfile(): void {
 		this.profile = this.notebookParams ? this.notebookParams.profile : undefined;
-		let profile: IConnectionProfile;
 		if (!this.profile) {
-			// Use connectionProfile passed in first
-			if (this._notebookParams.connectionProfileId !== undefined && this._notebookParams.connectionProfileId) {
-				profile = this.connectionManagementService.getConnectionProfileById(this._notebookParams.connectionProfileId);
-			} else {
-				// Second use global connection if possible
-				profile = TaskUtilities.getCurrentGlobalConnection(this.objectExplorerService, this.connectionManagementService, this.editorService);
-			}
+			// Second use global connection if possible
+			let profile: IConnectionProfile = TaskUtilities.getCurrentGlobalConnection(this.objectExplorerService, this.connectionManagementService, this.editorService);
+
 			// TODO use generic method to match kernel with valid connection that's compatible. For now, we only have 1
 			if (profile && profile.providerName) {
 				this.profile = profile;
 			} else {
 				// if not, try 1st active connection that matches our filter
-				let profiles = this.connectionManagementService.getActiveConnections();
-				if (profiles && profiles.length > 0) {
-					this.profile = profiles[0];
+				let activeProfiles = this.connectionManagementService.getActiveConnections();
+				if (activeProfiles && activeProfiles.length > 0) {
+					this.profile = activeProfiles[0];
 				}
 			}
 		}

src/sql/parts/notebook/notebookEditor.ts

@@ -95,7 +95,7 @@ export class NotebookEditor extends BaseEditor {
 			providerId: input.providerId ? input.providerId : DEFAULT_NOTEBOOK_PROVIDER,
 			providers: input.providers ? input.providers : [DEFAULT_NOTEBOOK_PROVIDER],
 			isTrusted: input.isTrusted,
-			connectionProfileId: input.connectionProfileId
+			profile: input.connectionProfile
 		};
 		bootstrapAngular(this.instantiationService,
 			NotebookModule,

src/sql/parts/notebook/notebookInput.ts

@@ -27,6 +27,7 @@ import { IUntitledEditorService } from 'vs/workbench/services/untitled/common/un
 import { notebookModeId } from 'sql/common/constants';
 import { ITextFileService, ISaveOptions } from 'vs/workbench/services/textfile/common/textfiles';
 import { LocalContentManager } from 'sql/workbench/services/notebook/node/localContentManager';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
 
@@ -128,7 +129,7 @@ export class NotebookInput extends EditorInput {
 	private _providerId: string;
 	private _providers: string[];
 	private _standardKernels: IStandardKernelWithProvider[];
-	private _connectionProfileId: string;
+	private _connectionProfile: IConnectionProfile;
 	private _defaultKernel: azdata.nb.IKernelSpec;
 	private _isTrusted: boolean = false;
 	public hasBootstrapped = false;
@@ -191,12 +192,12 @@ export class NotebookInput extends EditorInput {
 		this._isTrusted = value;
 	}
 
-	public set connectionProfileId(value: string) {
-		this._connectionProfileId = value;
+	public set connectionProfile(value: IConnectionProfile) {
+		this._connectionProfile = value;
 	}
 
-	public get connectionProfileId(): string {
-		return this._connectionProfileId;
+	public get connectionProfile(): IConnectionProfile {
+		return this._connectionProfile;
 	}
 
 	public get standardKernels(): IStandardKernelWithProvider[] {

src/sql/workbench/api/node/extHostNotebookDocumentsAndEditors.ts

@@ -167,7 +167,7 @@ export class ExtHostNotebookDocumentsAndEditors implements ExtHostNotebookDocume
 			options.preview = showOptions.preview;
 			options.position = showOptions.viewColumn;
 			options.providerId = showOptions.providerId;
-			options.connectionId = showOptions.connectionId;
+			options.connectionProfile = showOptions.connectionProfile;
 			options.defaultKernel = showOptions.defaultKernel;
 		}
 		let id = await this._proxy.$tryShowNotebookDocument(uri, options);

src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts

@@ -16,6 +16,7 @@ import { IInstantiationService } from 'vs/platform/instantiation/common/instanti
 import { ITextEditorOptions } from 'vs/platform/editor/common/editor';
 import { viewColumnToEditorGroup } from 'vs/workbench/api/shared/editor';
 import { Schemas } from 'vs/base/common/network';
+import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 
 import {
 	SqlMainContext, MainThreadNotebookDocumentsAndEditorsShape, SqlExtHostContext, ExtHostNotebookDocumentsAndEditorsShape,
@@ -28,6 +29,7 @@ import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHos
 import { disposed } from 'vs/base/common/errors';
 import { ICellModel, NotebookContentChange, INotebookModel } from 'sql/parts/notebook/models/modelInterfaces';
 import { NotebookChangeType, CellTypes } from 'sql/parts/notebook/models/contracts';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 class MainThreadNotebookEditor extends Disposable {
 	private _contentChangedEmitter = new Emitter<NotebookContentChange>();
@@ -291,7 +293,7 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@IEditorService private _editorService: IEditorService,
 		@IEditorGroupsService private _editorGroupService: IEditorGroupsService,
-		@INotebookService private readonly _notebookService: INotebookService
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
 	) {
 		super();
 		if (extHostContext) {
@@ -363,7 +365,7 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		let input = this._instantiationService.createInstance(NotebookInput, uri.fsPath, uri);
 		input.isTrusted = trusted;
 		input.defaultKernel = options.defaultKernel;
-		input.connectionProfileId = options.connectionId;
+		input.connectionProfile = new ConnectionProfile(this._capabilitiesService, options.connectionProfile);
 
 		let editor = await this._editorService.openEditor(input, editorOptions, viewColumnToEditorGroup(this._editorGroupService, options.position));
 		if (!editor) {

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -850,7 +850,7 @@ export interface INotebookShowOptions {
 	preserveFocus?: boolean;
 	preview?: boolean;
 	providerId?: string;
-	connectionId?: string;
+	connectionProfile?: azdata.IConnectionProfile;
 	defaultKernel?: azdata.nb.IKernelSpec;
 }
 

src/sql/workbench/services/commandLine/common/commandLineService.ts

@@ -3,6 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 'use strict';
+import * as azdata from 'azdata';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { ICommandLineProcessing } from 'sql/workbench/services/commandLine/common/commandLine';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
@@ -19,8 +20,10 @@ import { IWorkspaceConfigurationService } from 'vs/workbench/services/configurat
 import { ICommandService } from 'vs/platform/commands/common/commands';
 import { warn } from 'sql/base/common/log';
 import { ipcRenderer as ipc} from 'electron';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 
 export class CommandLineService implements ICommandLineProcessing {
+	public _serviceBrand: any;
 
 	constructor(
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
@@ -57,14 +60,13 @@ export class CommandLineService implements ICommandLineProcessing {
 		}
 	}
 
-	public _serviceBrand: any;
 	// We base our logic on the combination of (server, command) values.
 	// (serverName, commandName) => Connect object explorer and execute the command, passing the connection profile to the command. Do not load query editor.
 	// (null, commandName) => Launch the command with a null connection. If the command implementation needs a connection, it will need to create it.
 	// (serverName, null) => Connect object explorer and open a new query editor
 	// (null, null) => Prompt for a connection unless there are registered servers
-	public processCommandLine(args: ParsedArgs): Promise<void> {
-		let profile = undefined;
+	public async processCommandLine(args: ParsedArgs): Promise<void> {
+		let profile: IConnectionProfile = undefined;
 		let commandName = undefined;
 		if (args) {
 			if (this._commandService) {
@@ -72,66 +74,58 @@ export class CommandLineService implements ICommandLineProcessing {
 			}
 
 			if (args.server) {
-				profile = new ConnectionProfile(this._capabilitiesService, null);
-				// We want connection store to use any matching password it finds
-				profile.savePassword = true;
-				profile.providerName = Constants.mssqlProviderName;
-				profile.serverName = args.server;
-				profile.databaseName = args.database ? args.database : '';
-				profile.userName = args.user ? args.user : '';
-				profile.authenticationType = args.integrated ? 'Integrated' : 'SqlLogin';
-				profile.connectionName = '';
-				profile.setOptionValue('applicationName', Constants.applicationName);
-				profile.setOptionValue('databaseDisplayName', profile.databaseName);
-				profile.setOptionValue('groupId', profile.groupId);
+				profile = this.readProfileFromArgs(args);
 			}
 	    }
-		let self = this;
-		return new Promise<void>((resolve, reject) => {
-			let showConnectDialogOnStartup: boolean = self._configurationService.getValue('workbench.showConnectDialogOnStartup');
-			if (showConnectDialogOnStartup && !commandName && !profile && !self._connectionManagementService.hasRegisteredServers()) {
-				// prompt the user for a new connection on startup if no profiles are registered
-				self._connectionManagementService.showConnectionDialog()
-					.then(() => {
-						resolve();
-					},
-						error => {
-							reject(error);
-						});
-			} else if (profile) {
-				if (!commandName) {
-					self._connectionManagementService.connectIfNotConnected(profile, 'connection', true)
-						.then(() => {
-							TaskUtilities.newQuery(profile,
-								self._connectionManagementService,
-								self._queryEditorService,
-								self._objectExplorerService,
-								self._editorService)
-								.then(() => {
-									resolve();
-								}, error => {
-									// ignore query editor failing to open.
-									// the tests don't mock this out
-									warn('unable to open query editor ' + error);
-									resolve();
-								});
-						}, error => {
-							reject(error);
-						});
-				} else {
-					self._connectionManagementService.connectIfNotConnected(profile, 'connection', true)
-						.then(() => {
-							self._commandService.executeCommand(commandName, profile.id).then(() => resolve(), error => reject(error));
-						}, error => {
-							reject(error);
-						});
-				}
-			} else if (commandName) {
-				self._commandService.executeCommand(commandName).then(() => resolve(), error => reject(error));
+		let showConnectDialogOnStartup: boolean = this._configurationService.getValue('workbench.showConnectDialogOnStartup');
+		if (showConnectDialogOnStartup && !commandName && !profile && !this._connectionManagementService.hasRegisteredServers()) {
+			// prompt the user for a new connection on startup if no profiles are registered
+			await this._connectionManagementService.showConnectionDialog();
+			return;
+		}
+		let connectedContext: azdata.ConnectedContext = undefined;
+		if (profile) {
+			try {
+				await this._connectionManagementService.connectIfNotConnected(profile, 'connection', true);
+				// Before sending to extensions, we should a) serialize to IConnectionProfile or things will fail,
+				// and b) use the latest version of the profile from the service so most fields are filled in.
+				let updatedProfile = this._connectionManagementService.getConnectionProfileById(profile.id);
+				connectedContext = { connectionProfile: new ConnectionProfile(this._capabilitiesService, updatedProfile).toIConnectionProfile() };
+			} catch (err) {
+				warn('Failed to connect due to error' + err.message);
 			}
-			else {
-				resolve();
+		}
+		if (commandName) {
+			await this._commandService.executeCommand(commandName, connectedContext);
+		} else if (profile) {
+			// Default to showing new query
+			try {
+				await TaskUtilities.newQuery(profile,
+					this._connectionManagementService,
+					this._queryEditorService,
+					this._objectExplorerService,
+					this._editorService);
+			} catch (error) {
+				warn('unable to open query editor ' + error);
+				// Note: we are intentionally swallowing this error.
+				// In part this is to accommodate unit testing where we don't want to set up the query stack
 			}
-		});
+		}
+	}
+
+	private readProfileFromArgs(args: ParsedArgs) {
+		let profile = new ConnectionProfile(this._capabilitiesService, null);
+		// We want connection store to use any matching password it finds
+		profile.savePassword = true;
+		profile.providerName = Constants.mssqlProviderName;
+		profile.serverName = args.server;
+		profile.databaseName = args.database ? args.database : '';
+		profile.userName = args.user ? args.user : '';
+		profile.authenticationType = args.integrated ? 'Integrated' : 'SqlLogin';
+		profile.connectionName = '';
+		profile.setOptionValue('applicationName', Constants.applicationName);
+		profile.setOptionValue('databaseDisplayName', profile.databaseName);
+		profile.setOptionValue('groupId', profile.groupId);
+		return profile;
 	}
 }

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

@@ -95,7 +95,6 @@ export interface INotebookParams extends IBootstrapParams {
 	isTrusted: boolean;
 	profile?: IConnectionProfile;
 	modelFactory?: ModelFactory;
-	connectionProfileId?: string;
 }
 
 export interface INotebookEditor {

src/sqltest/parts/commandLine/commandLineService.test.ts

@@ -7,7 +7,7 @@
 
 import * as assert from 'assert';
 import * as TypeMoq from 'typemoq';
-
+import * as azdata from 'azdata';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { CommandLineService } from 'sql/workbench/services/commandLine/common/commandLineService';
 import { ParsedArgs } from 'vs/platform/environment/common/environment';
@@ -23,6 +23,8 @@ import { ICommandService } from 'vs/platform/commands/common/commands';
 import { TestCommandService } from 'vs/editor/test/browser/editorTestServices';
 import { WorkspaceConfigurationTestService } from 'sqltest/stubs/workspaceConfigurationTestService';
 import { IWorkspaceConfigurationService } from 'vs/workbench/services/configuration/common/configuration';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
+import { assertThrowsAsync } from 'sqltest/utils/testUtils';
 
 class TestParsedArgs implements ParsedArgs {
 	[arg: string]: any;
@@ -144,7 +146,7 @@ suite('commandLineService tests', () => {
 		}, error => { assert.fail(error, null, 'processCommandLine rejected ' + error); done(); });
 	});
 
-	test('processCommandLine does nothing if no server name and command name is provided and the configuration \'workbench.showConnectDialogOnStartup\' is set to false, even if registered servers exist', done => {
+	test('processCommandLine does nothing if no server name and command name is provided and the configuration \'workbench.showConnectDialogOnStartup\' is set to false, even if registered servers exist', async () => {
 		const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
 			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
 
@@ -155,12 +157,11 @@ suite('commandLineService tests', () => {
 		const configurationService = getConfigurationServiceMock(false);
 		let service = getCommandLineService(connectionManagementService.object, configurationService.object);
 
-		service.processCommandLine(new TestParsedArgs());
+		await service.processCommandLine(new TestParsedArgs());
 		connectionManagementService.verifyAll();
-		done();
 	});
 
-	test('processCommandLine does nothing if registered servers exist and no server name is provided', done => {
+	test('processCommandLine does nothing if registered servers exist and no server name is provided', async () => {
 		const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
 			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
 
@@ -171,13 +172,15 @@ suite('commandLineService tests', () => {
 			.verifiable(TypeMoq.Times.never());
 		const configurationService = getConfigurationServiceMock(true);
 		let service = getCommandLineService(connectionManagementService.object, configurationService.object);
-		service.processCommandLine(new TestParsedArgs()).then(() => {
+		try {
+			await service.processCommandLine(new TestParsedArgs());
 			connectionManagementService.verifyAll();
-			done();
-		}, error => { assert.fail(error, null, 'processCommandLine rejected ' + error); done(); });
+		} catch (error){
+			assert.fail(error, null, 'processCommandLine rejected ' + error);
+		}
 	});
 
-	test('processCommandLine opens a new connection if a server name is passed', done => {
+	test('processCommandLine opens a new connection if a server name is passed', async () => {
 		const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
 			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
 
@@ -186,20 +189,23 @@ suite('commandLineService tests', () => {
 		args.database = 'mydatabase';
 		connectionManagementService.setup((c) => c.showConnectionDialog()).verifiable(TypeMoq.Times.never());
 		connectionManagementService.setup(c => c.hasRegisteredServers()).returns(() => true).verifiable(TypeMoq.Times.atMostOnce());
-		connectionManagementService.setup(c => c.connectIfNotConnected(TypeMoq.It.isAny(), 'connection', true))
-			.returns(() => new Promise<string>((resolve, reject) => { resolve('unused'); }))
+		let originalProfile: IConnectionProfile = undefined;
+		connectionManagementService.setup(c => c.connectIfNotConnected(TypeMoq.It.is<ConnectionProfile>(p => p.serverName === 'myserver'), 'connection', true))
+			.returns((conn) => {
+				originalProfile = conn;
+				return Promise.resolve('unused');
+			})
 			.verifiable(TypeMoq.Times.once());
+		connectionManagementService.setup(c => c.getConnectionProfileById(TypeMoq.It.isAnyString())).returns(() => originalProfile);
 		const configurationService = getConfigurationServiceMock(true);
 		let service = getCommandLineService(connectionManagementService.object, configurationService.object, capabilitiesService);
-		service.processCommandLine(args).then(() => {
-			connectionManagementService.verifyAll();
-			done();
-		}, error => { assert.fail(error, null, 'processCommandLine rejected ' + error); done(); });
+		await service.processCommandLine(args);
+		connectionManagementService.verifyAll();
 	});
 
-	test('processCommandLine invokes a command without a profile parameter when no server is passed', done => {
+	test('processCommandLine invokes a command without a profile parameter when no server is passed', async () => {
 		const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
-			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Loose);
 		const commandService: TypeMoq.Mock<ICommandService> = TypeMoq.Mock.ofType<ICommandService>(TestCommandService);
 		const args: TestParsedArgs = new TestParsedArgs();
 
@@ -208,20 +214,23 @@ suite('commandLineService tests', () => {
 		connectionManagementService.setup(c => c.hasRegisteredServers()).returns(() => true).verifiable(TypeMoq.Times.atMostOnce());
 		connectionManagementService.setup(c => c.connectIfNotConnected(TypeMoq.It.isAny(), TypeMoq.It.isAny()))
 			.verifiable(TypeMoq.Times.never());
-		commandService.setup(c => c.executeCommand('mycommand'))
-			.returns(() => Promise.resolve())
+		let capturedArgs: any;
+		commandService.setup(c => c.executeCommand(TypeMoq.It.isAnyString(), undefined))
+			.returns((command, args) => {
+				capturedArgs = args;
+				return Promise.resolve();
+			})
 			.verifiable(TypeMoq.Times.once());
 		const configurationService = getConfigurationServiceMock(true);
 		let service = getCommandLineService(connectionManagementService.object, configurationService.object, capabilitiesService, commandService.object);
-		service.processCommandLine(args).then(() => {
-			connectionManagementService.verifyAll();
-			commandService.verifyAll();
-			done();
-		}, error => { assert.fail(error, null, 'processCommandLine rejected ' + error); done(); });
+		await service.processCommandLine(args);
+		connectionManagementService.verifyAll();
+		commandService.verifyAll();
+		should(capturedArgs).be.undefined();
 	});
 
 
-	test('processCommandLine invokes a command with a profile parameter when a server is passed', done => {
+	test('processCommandLine invokes a command with a profile parameter when a server is passed', async () => {
 
 		const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
 			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
@@ -232,22 +241,33 @@ suite('commandLineService tests', () => {
 		args.server = 'myserver';
 		connectionManagementService.setup((c) => c.showConnectionDialog()).verifiable(TypeMoq.Times.never());
 		connectionManagementService.setup(c => c.hasRegisteredServers()).returns(() => true).verifiable(TypeMoq.Times.atMostOnce());
+		let originalProfile: IConnectionProfile = undefined;
 		connectionManagementService.setup(c => c.connectIfNotConnected(TypeMoq.It.is<ConnectionProfile>(p => p.serverName === 'myserver'), 'connection', true))
-			.returns(() => new Promise<string>((resolve, reject) => { resolve('unused'); }))
+			.returns((conn) => {
+				originalProfile = conn;
+				return Promise.resolve('unused');
+			})
 			.verifiable(TypeMoq.Times.once());
-		commandService.setup(c => c.executeCommand('mycommand', TypeMoq.It.isAnyString()))
-			.returns(() => Promise.resolve())
+		connectionManagementService.setup(c => c.getConnectionProfileById(TypeMoq.It.isAnyString())).returns(() => originalProfile);
+
+		let actualProfile: azdata.ConnectedContext = undefined;
+		commandService.setup(c => c.executeCommand('mycommand', TypeMoq.It.isAny()))
+			.returns((cmdName, profile) => {
+				actualProfile = profile;
+				return Promise.resolve();
+			})
 			.verifiable(TypeMoq.Times.once());
 		const configurationService = getConfigurationServiceMock(true);
 		let service = getCommandLineService(connectionManagementService.object, configurationService.object, capabilitiesService, commandService.object);
-		service.processCommandLine(args).then(() => {
-			connectionManagementService.verifyAll();
-			commandService.verifyAll();
-			done();
-		}, error => { assert.fail(error, null, 'processCommandLine rejected ' + error); done(); });
+		await service.processCommandLine(args);
+		connectionManagementService.verifyAll();
+		commandService.verifyAll();
+		should(actualProfile).not.be.undefined();
+		should(actualProfile.connectionProfile.serverName).equal(args.server);
+
 	});
 
-	test('processCommandLine rejects unknown commands', done => {
+	test('processCommandLine rejects unknown commands', async () => {
 		const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
 			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
 		const commandService: TypeMoq.Mock<ICommandService> = TypeMoq.Mock.ofType<ICommandService>(TestCommandService);
@@ -260,13 +280,7 @@ suite('commandLineService tests', () => {
 			.verifiable(TypeMoq.Times.once());
 		const configurationService = getConfigurationServiceMock(true);
 		let service = getCommandLineService(connectionManagementService.object, configurationService.object, capabilitiesService, commandService.object);
-		service.processCommandLine(args).then(() => {
-			assert.fail(1, null, 'processCommandLine should reject when executeCommand errors out');
-			done();
-		}, error => {
-			assert.equal(error.message, 'myerror', 'unexpected error from processCommandLine ' + error);
-			done();
-		});
+		assertThrowsAsync(async () => await service.processCommandLine(args));
 	});
 });