Inital platform relayering (#6385)

* moving test files and inital refactoring

* relayer extension host code

* fix imports

* make insights work

* relayer dashboard

* relayer notebooks

* moveing more code around

* formatting

* accept angular as browser

* fix serializer

* add missing files

* remove declarations from extensions

* fix build errors

* more relayering

* change urls to relative to help code relayering

* remove layering to prep for merge

* fix hygiene errors

* fix hygiene errors

* fix tests

src/sql/workbench/test/electron-browser/api/extHostAccountManagement.test.ts

@@ -0,0 +1,460 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import * as azdata from 'azdata';
+import * as TypeMoq from 'typemoq';
+import { AccountProviderStub, TestAccountManagementService } from 'sql/platform/accounts/test/common/testAccountManagementService';
+import { ExtHostAccountManagement } from 'sql/workbench/api/common/extHostAccountManagement';
+import { TestRPCProtocol } from 'vs/workbench/test/electron-browser/api/testRPCProtocol';
+import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
+import { IRPCProtocol } from 'vs/workbench/services/extensions/common/proxyIdentifier';
+import { SqlMainContext } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { MainThreadAccountManagement } from 'sql/workbench/api/browser/mainThreadAccountManagement';
+import { IAccountManagementService, AzureResource } from 'sql/platform/accounts/common/interfaces';
+import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
+
+const IRPCProtocol = createDecorator<IRPCProtocol>('rpcProtocol');
+
+// SUITE STATE /////////////////////////////////////////////////////////////
+let instantiationService: TestInstantiationService;
+let mockAccountMetadata: azdata.AccountProviderMetadata;
+let mockAccount: azdata.Account;
+let threadService: TestRPCProtocol;
+
+// TESTS ///////////////////////////////////////////////////////////////////
+suite('ExtHostAccountManagement', () => {
+	suiteSetup(() => {
+		threadService = new TestRPCProtocol();
+		const accountMgmtStub = new TestAccountManagementService();
+
+		instantiationService = new TestInstantiationService();
+		instantiationService.stub(IRPCProtocol, threadService);
+		instantiationService.stub(IAccountManagementService, accountMgmtStub);
+
+		const accountMgmtService = instantiationService.createInstance(MainThreadAccountManagement, undefined);
+		threadService.set(SqlMainContext.MainThreadAccountManagement, accountMgmtService);
+
+		mockAccountMetadata = {
+			args: {},
+			displayName: 'Test Account Provider',
+			id: 'test_account_provider',
+			settings: {}
+		};
+		mockAccount = {
+			key: {
+				providerId: mockAccountMetadata.id,
+				providerArgs: {},
+				accountId: 'test_account'
+			},
+			properties: {},
+			displayInfo: {
+				displayName: 'Test Account',
+				userId: 'user@email.com',
+				contextualDisplayName: 'Test Kind Of Account',
+				accountType: 'test'
+			},
+			isStale: false
+		};
+	});
+
+	test('Constructor', () => {
+		// If: I construct a new extension host account management
+		let extHost = new ExtHostAccountManagement(threadService);
+
+		// Then: There shouldn't be any account providers registered
+		assert.equal(extHost.getProviderCount(), 0);
+	});
+
+	// REGISTER TESTS //////////////////////////////////////////////////////
+	test('Register Account Provider - Success', () => {
+		// Setup: Create an extension host account management
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+
+		// If: I register a mock account provider
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// Then: The account provider should be registered
+		assert.equal(extHost.getProviderCount(), 1);
+	});
+
+	test('Register Account Provider - Account Provider Already Registered', () => {
+		// Setup: Create an extension host account management and register an account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I register an account provider again
+		// Then
+		// ... It should throw
+		assert.throws(() => {
+			extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+		});
+
+		// ... There should only be one account provider
+		assert.equal(extHost.getProviderCount(), 1);
+	});
+
+	// TODO: Test for unregistering a provider
+
+	// CLEAR TESTS /////////////////////////////////////////////////////////
+	test('Clear - Success', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I clear an account
+		extHost.$clear(0, mockAccount.key)
+			.then(() => {
+				// Then: The call should have been passed to the provider
+				mockProvider.verify(
+					(obj) => obj.clear(TypeMoq.It.isValue(mockAccount.key)),
+					TypeMoq.Times.once()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	test('Clear - Handle does not exist', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I clear an account for a handle that doesn't exist
+		// Then: It should fail
+		extHost.$clear(1, mockAccount.key)
+			.then(() => done('Clear succeeded when it should have failed'))
+			.then(null, () => {
+				// The provider's clear should not have been called
+				mockProvider.verify(
+					(obj) => obj.clear(TypeMoq.It.isAny()),
+					TypeMoq.Times.never()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	// INITIALIZE TESTS ////////////////////////////////////////////////////
+	test('Initialize - Success', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I initialize the provider
+		extHost.$initialize(0, [mockAccount])
+			.then(() => {
+				// Then: The call should have been passed to the provider
+				mockProvider.verify(
+					(obj) => obj.initialize(TypeMoq.It.isValue([mockAccount])),
+					TypeMoq.Times.once()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	test('Initialize - Handle does not exist', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I initialize for a handle that doesn't exist
+		// Then: It should fail
+		extHost.$initialize(1, [mockAccount])
+			.then(() => done('Initialize succeeded when it should have failed'))
+			.then(null, () => {
+				// The provider's clear should not have been called
+				mockProvider.verify(
+					(obj) => obj.initialize(TypeMoq.It.isAny()),
+					TypeMoq.Times.never()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	// PROMPT TESTS ////////////////////////////////////////////////////////
+	test('Prompt - Success', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I prompt for an account
+		extHost.$prompt(0)
+			.then(() => {
+				// Then: The call should have been passed to the provider
+				mockProvider.verify(
+					(obj) => obj.prompt(),
+					TypeMoq.Times.once()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	test('Prompt - Handle does not exist', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I prompt with a handle that doesn't exist
+		// Then: It should fail
+		extHost.$prompt(1)
+			.then(() => done('Prompt succeeded when it should have failed'))
+			.then(null, () => {
+				// The provider's clear should not have been called
+				mockProvider.verify(
+					(obj) => obj.prompt(),
+					TypeMoq.Times.never()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	// REFRESH TESTS ///////////////////////////////////////////////////////
+	test('Refresh - Success', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I refresh an account
+		extHost.$refresh(0, mockAccount)
+			.then(() => {
+				// Then: The call should have been passed to the provider
+				mockProvider.verify(
+					(obj) => obj.refresh(TypeMoq.It.isValue(mockAccount)),
+					TypeMoq.Times.once()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	test('Refresh - Handle does not exist', (done) => {
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		let mockProvider = getMockAccountProvider();
+		extHost.$registerAccountProvider(mockAccountMetadata, mockProvider.object);
+
+		// If: I refresh an account for a handle that doesn't exist
+		// Then: It should fail
+		extHost.$refresh(1, mockAccount)
+			.then(() => done('Refresh succeeded when it should have failed'))
+			.then(null, () => {
+				// The provider's clear should not have been called
+				mockProvider.verify(
+					(obj) => obj.refresh(TypeMoq.It.isAny()),
+					TypeMoq.Times.never()
+				);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	// GETALLACCOUNTS TESTS ///////////////////////////////////////////////////////
+	test('GetAllAccounts - Success', (done) => {
+		let mockAccountProviderMetadata = {
+			id: 'azure',
+			displayName: 'Azure'
+		};
+
+		let mockAccount1 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_1'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Microsoft Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 1',
+				userId: 'user@email.com'
+			},
+			properties: [],
+			isStale: false
+		};
+		let mockAccount2 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_2'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Work/School Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 2',
+				userId: 'user@email.com'
+			},
+			properties: [],
+			isStale: false
+		};
+		let mockAccounts = [mockAccount1, mockAccount2];
+
+		let expectedAccounts = [mockAccount1, mockAccount2];
+
+		let mockAccountManagementService = getMockAccountManagementService(mockAccounts);
+		instantiationService.stub(IAccountManagementService, mockAccountManagementService.object);
+		let accountManagementService = instantiationService.createInstance(MainThreadAccountManagement, undefined);
+		threadService.set(SqlMainContext.MainThreadAccountManagement, accountManagementService);
+
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		extHost.$registerAccountProvider(mockAccountProviderMetadata, new AccountProviderStub());
+
+		// If: I get all accounts
+		extHost.$getAllAccounts()
+			.then((accounts) => {
+				// Then: The call should have been passed to the account management service
+				mockAccountManagementService.verify(
+					(obj) => obj.getAccountsForProvider(TypeMoq.It.isAny()),
+					TypeMoq.Times.once()
+				);
+
+				assert.ok(Array.isArray(accounts));
+				assert.equal(accounts.length, expectedAccounts.length);
+				assert.deepStrictEqual(accounts, expectedAccounts);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	test('GetAllAccounts - No account providers', (done) => {
+		// Setup: Create ext host account management with no registered account providers
+		let extHost = new ExtHostAccountManagement(threadService);
+
+		// If: I get all accounts
+		// Then: It should throw
+		assert.throws(
+			() => extHost.$getAllAccounts(),
+			(error) => {
+				return error.message === 'No account providers registered.';
+			});
+		done();
+	});
+
+	test('GetSecurityToken - Success', (done) => {
+		let mockAccountProviderMetadata = {
+			id: 'azure',
+			displayName: 'Azure'
+		};
+
+		let mockAccount1 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_1'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Microsoft Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 1',
+				userId: 'user@email.com'
+			},
+			properties: [],
+			isStale: false
+		};
+		let mockAccounts = [mockAccount1];
+
+		let mockAccountManagementService = getMockAccountManagementService(mockAccounts);
+		instantiationService.stub(IAccountManagementService, mockAccountManagementService.object);
+		let accountManagementService = instantiationService.createInstance(MainThreadAccountManagement, undefined);
+		threadService.set(SqlMainContext.MainThreadAccountManagement, accountManagementService);
+
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		extHost.$registerAccountProvider(mockAccountProviderMetadata, new AccountProviderStub());
+
+		extHost.$getAllAccounts()
+			.then((accounts) => {
+				// If: I get security token it will not throw
+				return extHost.$getSecurityToken(mockAccount1, AzureResource.ResourceManagement);
+			}
+			).then(() => done(), (err) => done(new Error(err)));
+	});
+
+	test('GetSecurityToken - Account not found', (done) => {
+		let mockAccountProviderMetadata = {
+			id: 'azure',
+			displayName: 'Azure'
+		};
+
+		let mockAccount1 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_1'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Microsoft Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 1',
+				userId: 'user@email.com'
+			},
+			properties: [],
+			isStale: false
+		};
+		let mockAccounts = [mockAccount1];
+
+		let mockAccountManagementService = getMockAccountManagementService(mockAccounts);
+		instantiationService.stub(IAccountManagementService, mockAccountManagementService.object);
+		let accountManagementService = instantiationService.createInstance(MainThreadAccountManagement, undefined);
+		threadService.set(SqlMainContext.MainThreadAccountManagement, accountManagementService);
+
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		extHost.$registerAccountProvider(mockAccountProviderMetadata, new AccountProviderStub());
+
+		let mockAccount2 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_2'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Work/School Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 2',
+				userId: 'user@email.com'
+			},
+			properties: [],
+			isStale: false
+		};
+
+		extHost.$getAllAccounts()
+			.then(accounts => {
+				return extHost.$getSecurityToken(mockAccount2, AzureResource.ResourceManagement);
+			})
+			.then((noError) => {
+				done(new Error('Expected getSecurityToken to throw'));
+			}, (err) => {
+				// Expected error caught
+				done();
+			});
+	});
+});
+
+function getMockAccountProvider(): TypeMoq.Mock<azdata.AccountProvider> {
+	let mock = TypeMoq.Mock.ofType<azdata.AccountProvider>(AccountProviderStub);
+	mock.setup((obj) => obj.clear(TypeMoq.It.isValue(mockAccount.key)))
+		.returns(() => Promise.resolve(undefined));
+	mock.setup((obj) => obj.refresh(TypeMoq.It.isValue(mockAccount)))
+		.returns(() => Promise.resolve(undefined));
+	mock.setup((obj) => obj.initialize(TypeMoq.It.isValue([mockAccount])))
+		.returns(() => Promise.resolve(undefined));
+	mock.setup((obj) => obj.prompt())
+		.returns(() => Promise.resolve(undefined));
+
+	return mock;
+}
+
+function getMockAccountManagementService(accounts: azdata.Account[]): TypeMoq.Mock<TestAccountManagementService> {
+	let mockAccountManagementService = TypeMoq.Mock.ofType(TestAccountManagementService);
+
+	mockAccountManagementService.setup(x => x.getAccountsForProvider(TypeMoq.It.isAny()))
+		.returns(() => Promise.resolve(accounts));
+	mockAccountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isValue(accounts[0]), TypeMoq.It.isAny()))
+		.returns(() => Promise.resolve({}));
+	mockAccountManagementService.setup(x => x.updateAccountListEvent)
+		.returns(() => () => { return undefined; });
+
+	return mockAccountManagementService;
+}

src/sql/workbench/test/electron-browser/api/extHostBackgroundTaskManagement.test.ts

@@ -0,0 +1,106 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import * as assert from 'assert';
+import { Mock, It, Times } from 'typemoq';
+import { ExtHostBackgroundTaskManagement, TaskStatus } from 'sql/workbench/api/common/extHostBackgroundTaskManagement';
+import { MainThreadBackgroundTaskManagementShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
+
+
+suite('ExtHostBackgroundTaskManagement Tests', () => {
+	let extHostBackgroundTaskManagement: ExtHostBackgroundTaskManagement;
+	let mockProxy: Mock<MainThreadBackgroundTaskManagementShape>;
+	let nothing: void;
+	let operationId = 'operation is';
+	setup(() => {
+		mockProxy = Mock.ofInstance(<MainThreadBackgroundTaskManagementShape>{
+
+			$registerTask: (taskInfo: azdata.TaskInfo) => nothing,
+			$updateTask: (taskProgressInfo: azdata.TaskProgressInfo) => nothing
+		});
+		let mainContext = <IMainContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+
+		mockProxy.setup(x => x.$registerTask(It.isAny())).callback(() => {
+			extHostBackgroundTaskManagement.$onTaskRegistered(operationId);
+		});
+		extHostBackgroundTaskManagement = new ExtHostBackgroundTaskManagement(mainContext);
+	});
+
+	test('RegisterTask should successfully create background task and update status', () => {
+		let operationInfo: azdata.BackgroundOperationInfo = {
+			connection: undefined,
+			description: 'description',
+			displayName: 'displayName',
+			isCancelable: true,
+			operation: (op: azdata.BackgroundOperation) => { op.updateStatus(TaskStatus.Succeeded); },
+			operationId: operationId
+		};
+		extHostBackgroundTaskManagement.$registerTask(operationInfo);
+		mockProxy.verify(x => x.$registerTask(It.is(
+			t => t.name === operationInfo.displayName &&
+				t.description === operationInfo.description &&
+				t.taskId === operationId &&
+				t.isCancelable === operationInfo.isCancelable &&
+				t.providerName === undefined
+		)), Times.once());
+		mockProxy.verify(x => x.$updateTask(It.is(t => t.status === TaskStatus.Succeeded)), Times.once());
+		extHostBackgroundTaskManagement.$removeTask(operationId);
+	});
+
+	test('Canceling the task should notify the extension', () => {
+		let operationInfo: azdata.BackgroundOperationInfo = {
+			connection: undefined,
+			description: 'description',
+			displayName: 'displayName',
+			isCancelable: true,
+			operation: (op: azdata.BackgroundOperation) => {
+				op.onCanceled(() => {
+					op.updateStatus(TaskStatus.Canceled);
+				});
+			},
+			operationId: operationId
+		};
+		extHostBackgroundTaskManagement.$registerTask(operationInfo);
+		extHostBackgroundTaskManagement.$onTaskCanceled(operationId);
+
+		mockProxy.verify(x => x.$updateTask(It.is(t => t.status === TaskStatus.Canceled)), Times.once());
+		extHostBackgroundTaskManagement.$removeTask(operationId);
+	});
+
+	test('RegisterTask should assign unique id to the operation is not assigned', () => {
+		let operationInfo: azdata.BackgroundOperationInfo = {
+			connection: undefined,
+			description: 'description',
+			displayName: 'displayName',
+			isCancelable: true,
+			operation: (op: azdata.BackgroundOperation) => { op.updateStatus(TaskStatus.Succeeded); },
+			operationId: undefined
+		};
+		extHostBackgroundTaskManagement.$registerTask(operationInfo);
+		mockProxy.verify(x => x.$registerTask(It.is(t => t.taskId !== undefined)), Times.once());
+
+		extHostBackgroundTaskManagement.$removeTask(operationId);
+	});
+
+	test('RegisterTask should fail given id of an existing operation', () => {
+		let operationInfo: azdata.BackgroundOperationInfo = {
+			connection: undefined,
+			description: 'description',
+			displayName: 'displayName',
+			isCancelable: true,
+			operation: (op: azdata.BackgroundOperation) => { op.updateStatus(TaskStatus.Succeeded); },
+			operationId: operationId
+		};
+		extHostBackgroundTaskManagement.$registerTask(operationInfo);
+		mockProxy.verify(x => x.$registerTask(It.is(t => t.taskId === operationId)), Times.once());
+		assert.throws(() => extHostBackgroundTaskManagement.$registerTask(operationInfo));
+
+		extHostBackgroundTaskManagement.$removeTask(operationId);
+	});
+});

src/sql/workbench/test/electron-browser/api/extHostCredentialManagement.test.ts

@@ -0,0 +1,131 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import { TestRPCProtocol } from 'vs/workbench/test/electron-browser/api/testRPCProtocol';
+import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
+import { ExtHostCredentialManagement } from 'sql/workbench/api/common/extHostCredentialManagement';
+import { SqlMainContext } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { IRPCProtocol } from 'vs/workbench/services/extensions/common/proxyIdentifier';
+import { MainThreadCredentialManagement } from 'sql/workbench/api/browser/mainThreadCredentialManagement';
+import { ICredentialsService } from 'sql/platform/credentials/common/credentialsService';
+import { Credential, CredentialProvider } from 'azdata';
+import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
+import { TestCredentialsService, TestCredentialsProvider } from 'sql/platform/credentials/test/common/testCredentialsService';
+
+const IRPCProtocol = createDecorator<IRPCProtocol>('rpcProtocol');
+
+// SUITE STATE /////////////////////////////////////////////////////////////
+let credentialServiceStub: TestCredentialsService;
+let instantiationService: TestInstantiationService;
+let threadService: TestRPCProtocol;
+
+// TESTS ///////////////////////////////////////////////////////////////////
+suite('ExtHostCredentialManagement', () => {
+	suiteSetup(() => {
+		threadService = new TestRPCProtocol();
+		credentialServiceStub = new TestCredentialsService();
+
+		instantiationService = new TestInstantiationService();
+		instantiationService.stub(IRPCProtocol, threadService);
+		instantiationService.stub(ICredentialsService, credentialServiceStub);
+
+		const credentialService = instantiationService.createInstance(MainThreadCredentialManagement, undefined);
+		threadService.set(SqlMainContext.MainThreadCredentialManagement, credentialService);
+	});
+
+	test('Construct ExtHostCredentialManagement', () => {
+		// If: I construct an credential management extension host
+		let extHost = new ExtHostCredentialManagement(threadService);
+
+		// Then: The extension host should not have any providers registered
+		assert.equal(extHost.getProviderCount(), 0);
+	});
+
+	test('Register Credential Provider', () => {
+		// Setup: Create a mock credential provider
+		let extHost = new ExtHostCredentialManagement(threadService);
+		let mockCredentialProvider = new TestCredentialsProvider();
+
+		// If: I register the credential provider with the extension host
+		extHost.$registerCredentialProvider(mockCredentialProvider);
+
+		// Then: There should be one provider registered
+		assert.equal(extHost.getProviderCount(), 1);
+	});
+
+	test('Get Credential Provider - Success', (done) => {
+		// Setup: Register a mock credential provider
+		let extHost = new ExtHostCredentialManagement(threadService);
+		let mockCredentialProvider = new TestCredentialsProvider();
+		extHost.$registerCredentialProvider(mockCredentialProvider);
+
+		// If: I get the credential provider
+		let namespaceId = 'test_namespace';
+		let credentialId = 'test_id';
+		let credential = 'test_credential';
+		let expectedCredentialId = `${namespaceId}|${credentialId}`;
+		let credProvider: CredentialProvider;
+		extHost.$getCredentialProvider(namespaceId)
+			.then((provider) => {
+				// Then: There should still only be one provider registered
+				assert.equal(extHost.getProviderCount(), 1);
+				credProvider = provider;
+			})
+			.then(() => {
+				// If: I write a credential
+				return credProvider.saveCredential(credentialId, credential);
+			})
+			.then(() => {
+				// Then: The credential should have been stored with its namespace
+				assert.notStrictEqual(mockCredentialProvider.storedCredentials[expectedCredentialId], undefined);
+				assert.equal(mockCredentialProvider.storedCredentials[expectedCredentialId].credentialId, expectedCredentialId);
+				assert.equal(mockCredentialProvider.storedCredentials[expectedCredentialId].password, credential);
+			})
+			.then(() => {
+				// If: I read a credential
+				return credProvider.readCredential(credentialId);
+			})
+			.then((returnedCredential: Credential) => {
+				// Then: The credential ID should be namespaced
+				assert.equal(returnedCredential.credentialId, expectedCredentialId);
+				assert.equal(returnedCredential.password, credential);
+			})
+			.then(() => {
+				// If: I delete a credential
+				return credProvider.deleteCredential(credentialId);
+			})
+			.then(() => {
+				// Then: The credential with its namespace should no longer exist
+				assert.strictEqual(mockCredentialProvider.storedCredentials[expectedCredentialId], undefined);
+			})
+			.then(() => done(), (err) => done(err));
+	});
+
+	test('Get Credential Provider - No Namespace', (done) => {
+		// Setup: Register a mock credential provider
+		let extHost = new ExtHostCredentialManagement(threadService);
+		let mockCredentialProvider = new TestCredentialsProvider();
+		extHost.$registerCredentialProvider(mockCredentialProvider);
+
+		// If: I get a credential provider with an invalid namespace ID
+		// Then: I should get an error
+		extHost.$getCredentialProvider(undefined)
+			.then(
+				() => { done('Provider was returned from undefined'); },
+				() => { /* Swallow error, this is success path */ }
+			)
+			.then(() => { return extHost.$getCredentialProvider(null); })
+			.then(
+				() => { done('Provider was returned from null'); },
+				() => { /* Swallow error, this is success path */ }
+			)
+			.then(() => { return extHost.$getCredentialProvider(''); })
+			.then(
+				() => { done('Provider was returned from \'\''); },
+				() => { done(); }
+			);
+	});
+});

src/sql/workbench/test/electron-browser/api/extHostDataProtocol.test.ts

@@ -0,0 +1,61 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import { Mock } from 'typemoq';
+import * as azdata from 'azdata';
+import { ExtHostDataProtocol } from 'sql/workbench/api/common/extHostDataProtocol';
+import { DataProviderType } from 'sql/workbench/api/common/sqlExtHostTypes';
+
+suite('ExtHostDataProtocol', () => {
+
+	let extHostDataProtocol: ExtHostDataProtocol;
+
+	setup(() => {
+		extHostDataProtocol = new ExtHostDataProtocol({
+			getProxy: identifier => {
+				return {
+					$registerMetadataProvider: (providerId, handle) => Promise.resolve(),
+					$registerConnectionProvider: (providerId, handle) => Promise.resolve()
+				} as any;
+			}
+		} as any);
+	});
+
+	test('Providers are exposed to other extensions', () => {
+		let extension1Id = 'provider1';
+		let extension1MetadataMock = Mock.ofInstance({
+			getMetadata: () => undefined,
+			getDatabases: () => undefined,
+			getTableInfo: () => undefined,
+			getViewInfo: () => undefined,
+			providerId: extension1Id
+		} as azdata.MetadataProvider);
+
+		let extension2Id = 'provider2';
+		let extension2MetadataMock = Mock.ofInstance({
+			getMetadata: () => undefined,
+			getDatabases: () => undefined,
+			getTableInfo: () => undefined,
+			getViewInfo: () => undefined,
+			providerId: extension2Id
+		} as azdata.MetadataProvider);
+
+		// If I register both providers and then get them using the getProvider API
+		extHostDataProtocol.$registerMetadataProvider(extension1MetadataMock.object);
+		extHostDataProtocol.$registerMetadataProvider(extension2MetadataMock.object);
+		extHostDataProtocol.$registerConnectionProvider({} as azdata.ConnectionProvider);
+		let retrievedProvider1 = extHostDataProtocol.getProvider<azdata.MetadataProvider>(extension1Id, DataProviderType.MetadataProvider);
+		let retrievedProvider2 = extHostDataProtocol.getProvider<azdata.MetadataProvider>(extension2Id, DataProviderType.MetadataProvider);
+		let allProviders = extHostDataProtocol.getProvidersByType<azdata.MetadataProvider>(DataProviderType.MetadataProvider);
+
+		// Then each provider was retrieved successfully
+		assert.equal(retrievedProvider1, extension1MetadataMock.object, 'Expected metadata provider was not retrieved for extension 1');
+		assert.equal(retrievedProvider2, extension2MetadataMock.object, 'Expected metadata provider was not retrieved for extension 2');
+		assert.equal(allProviders.length, 2, 'All metadata providers had unexpected length');
+		assert.equal(allProviders.some(provider => provider === extension1MetadataMock.object), true, 'All metadata providers did not include extension 1 metadata provider');
+		assert.equal(allProviders.some(provider => provider === extension2MetadataMock.object), true, 'All metadata providers did not include extension 2 metadata provider');
+	});
+});

src/sql/workbench/test/electron-browser/api/extHostModelView.test.ts

@@ -0,0 +1,381 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import { Mock, It, Times, MockBehavior } from 'typemoq';
+import * as azdata from 'azdata';
+import { ExtHostModelView } from 'sql/workbench/api/common/extHostModelView';
+import { MainThreadModelViewShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
+import { IComponentShape, IItemConfig, ComponentEventType, IComponentEventArgs, ModelComponentTypes } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { TitledFormItemLayout } from 'sql/workbench/browser/modelComponents/formContainer.component';
+
+interface InternalItemConfig {
+	toIItemConfig(): IItemConfig;
+}
+interface IWithItemConfig {
+	itemConfigs?: InternalItemConfig[];
+}
+
+suite('ExtHostModelView Validation Tests', () => {
+	let extHostModelView: ExtHostModelView;
+	let mockProxy: Mock<MainThreadModelViewShape>;
+	let modelView: azdata.ModelView;
+	let inputBox: azdata.InputBoxComponent;
+	let validText = 'valid';
+	let widgetId = 'widget_id';
+	let handle = 1;
+	// let viewInitialized: Deferred<void>;
+	let initializedModels: IComponentShape[];
+
+	setup(done => {
+		// Set up the MainThreadModelViewShape proxy
+		mockProxy = Mock.ofInstance(<MainThreadModelViewShape>{
+			$registerProvider: (id: string) => undefined,
+			$initializeModel: (handle: number, rootComponent: IComponentShape) => undefined,
+			$clearContainer: (handle: number, componentId: string) => undefined,
+			$addToContainer: (handle: number, containerId: string, item: IItemConfig) => undefined,
+			$removeFromContainer: (handle: number, containerId: string, item: IItemConfig) => undefined,
+			$setLayout: (handle: number, componentId: string, layout: any) => undefined,
+			$setProperties: (handle: number, componentId: string, properties: { [key: string]: any }) => undefined,
+			$registerEvent: (handle: number, componentId: string) => undefined,
+			dispose: () => undefined,
+			$validate: (handle: number, componentId: string) => undefined
+		}, MockBehavior.Loose);
+		let mainContext = <IMainContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$registerEvent(It.isAny(), It.isAny())).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$setProperties(It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+
+		// Register a model view of an input box and drop down box inside a form container inside a flex container
+		extHostModelView = new ExtHostModelView(mainContext, undefined);
+		extHostModelView.$registerProvider(widgetId, async view => {
+			modelView = view;
+			inputBox = view.modelBuilder.inputBox()
+				.withValidation(component => component.value === validText)
+				.component();
+			let dropDownBox = view.modelBuilder.dropDown().component();
+			let formContainer = view.modelBuilder.formContainer()
+				.withItems([inputBox, dropDownBox])
+				.component();
+			let flexContainer = view.modelBuilder.flexContainer()
+				.withItems([formContainer])
+				.component();
+			await view.initializeModel(flexContainer);
+			done();
+		}, undefined);
+
+		extHostModelView.$registerWidget(handle, widgetId, undefined, undefined);
+	});
+
+	test('The custom validation output of a component gets set when it is initialized', done => {
+		extHostModelView.$runCustomValidations(handle, inputBox.id).then(valid => {
+			try {
+				assert.equal(valid, false, 'Empty input box did not validate as false');
+				done();
+			} catch (err) {
+				done(err);
+			}
+		}, err => done(err));
+	});
+
+	test('The custom validation output of a component changes if its value changes', done => {
+		inputBox.value = validText;
+		extHostModelView.$runCustomValidations(handle, inputBox.id).then(valid => {
+			try {
+				assert.equal(valid, true, 'Valid input box did not validate as valid');
+				done();
+			} catch (err) {
+				done(err);
+			}
+		}, err => done(err));
+	});
+
+	test('The custom validation output of a component changes after a PropertiesChanged event', done => {
+		extHostModelView.$handleEvent(handle, inputBox.id, {
+			args: {
+				'value': validText
+			},
+			eventType: ComponentEventType.PropertiesChanged
+		} as IComponentEventArgs);
+		extHostModelView.$runCustomValidations(handle, inputBox.id).then(valid => {
+			try {
+				assert.equal(valid, true, 'Valid input box did not validate as valid after PropertiesChanged event');
+				done();
+			} catch (err) {
+				done(err);
+			}
+		}, err => done(err));
+	});
+
+	test('The validity of a component is set by main thread validationChanged events', () => {
+		assert.equal(inputBox.valid, true, 'Component validity is true by default');
+		extHostModelView.$handleEvent(handle, inputBox.id, {
+			eventType: ComponentEventType.validityChanged,
+			args: false
+		});
+		assert.equal(inputBox.valid, false, 'Input box did not update validity to false based on the validityChanged event');
+		extHostModelView.$handleEvent(handle, inputBox.id, {
+			eventType: ComponentEventType.validityChanged,
+			args: true
+		});
+		assert.equal(inputBox.valid, true, 'Input box did not update validity to true based on the validityChanged event');
+	});
+
+	test('Main thread validityChanged events cause component to fire validity changed events', () => {
+		let validityFromEvent: boolean = undefined;
+		inputBox.onValidityChanged(valid => validityFromEvent = valid);
+		extHostModelView.$handleEvent(handle, inputBox.id, {
+			eventType: ComponentEventType.validityChanged,
+			args: false
+		});
+		assert.equal(validityFromEvent, false, 'Main thread validityChanged event did not cause component to fire its own event');
+	});
+
+	test('Setting a form component as required initializes the model with the component required', () => {
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+
+		// Set up the input component with required initially set to false
+		let inputComponent = modelView.modelBuilder.inputBox().component();
+		inputComponent.required = false;
+
+		// If I build a form that sets the input component as required
+		let inputFormComponent: azdata.FormComponent = {
+			component: inputComponent,
+			title: 'test_input',
+			required: true
+		};
+		let requiredFormContainer = modelView.modelBuilder.formContainer().withFormItems([inputFormComponent]).component();
+		modelView.initializeModel(requiredFormContainer);
+
+		// Then the input component is sent to the main thread with required set to true
+		mockProxy.verify(x => x.$initializeModel(It.isAny(), It.is(rootComponent => {
+			return rootComponent.itemConfigs.length === 1 && rootComponent.itemConfigs[0].componentShape.id === inputComponent.id && rootComponent.itemConfigs[0].componentShape.properties['required'] === true;
+		})), Times.once());
+	});
+
+	test('Form component groups are handled correctly by adding each item in the group and a label to the form', () => {
+		// Set up the mock proxy to save the component that gets initialized so that it can be verified
+		let rootComponent: IComponentShape;
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).callback((handle, componentShape) => rootComponent = componentShape);
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+
+		// Set up the form with a top level component and a group
+		let topLevelList = modelView.modelBuilder.listBox().component();
+		let groupInput = modelView.modelBuilder.inputBox().component();
+		let groupDropdown = modelView.modelBuilder.dropDown().component();
+
+		let topLevelInputFormComponent: azdata.FormComponent = { component: topLevelList, title: 'top_level_input' };
+		let groupInputFormComponent: azdata.FormComponent = { component: groupInput, title: 'group_input' };
+		let groupDropdownFormComponent: azdata.FormComponent = { component: groupDropdown, title: 'group_dropdown' };
+
+		let groupTitle = 'group_title';
+
+		// Give the group a default layout and add one just for the input component too
+		let defaultLayout: azdata.FormItemLayout = {
+			horizontal: true
+		};
+		let groupInputLayout: azdata.FormItemLayout = {
+			horizontal: false
+		};
+
+		// If I build a form that has a group with a default layout where one item in the group has its own layout
+		let formContainer = modelView.modelBuilder.formContainer().withFormItems([
+			topLevelInputFormComponent,
+			{
+				components: [
+					Object.assign(groupInputFormComponent, { layout: groupInputLayout }),
+					groupDropdownFormComponent
+				],
+				title: groupTitle
+			}
+		], defaultLayout).component();
+		modelView.initializeModel(formContainer);
+
+		// Then all the items plus a group label are added and have the correct layouts
+		assert.equal(rootComponent.itemConfigs.length, 4);
+		let listBoxConfig = rootComponent.itemConfigs[0];
+		let groupLabelConfig = rootComponent.itemConfigs[1];
+		let inputBoxConfig = rootComponent.itemConfigs[2];
+		let dropdownConfig = rootComponent.itemConfigs[3];
+
+		// Verify that the correct items were added
+		assert.equal(listBoxConfig.componentShape.type, ModelComponentTypes.ListBox);
+		assert.equal(groupLabelConfig.componentShape.type, ModelComponentTypes.Text);
+		assert.equal(inputBoxConfig.componentShape.type, ModelComponentTypes.InputBox);
+		assert.equal(dropdownConfig.componentShape.type, ModelComponentTypes.DropDown);
+
+		// Verify that the group title was set up correctly
+		assert.equal(groupLabelConfig.componentShape.properties['value'], groupTitle);
+		assert.equal((groupLabelConfig.config as TitledFormItemLayout).isGroupLabel, true);
+
+		// Verify that the components' layouts are correct
+		assert.equal((listBoxConfig.config as azdata.FormItemLayout).horizontal, defaultLayout.horizontal);
+		assert.equal((inputBoxConfig.config as azdata.FormItemLayout).horizontal, groupInputLayout.horizontal);
+		assert.equal((dropdownConfig.config as azdata.FormItemLayout).horizontal, defaultLayout.horizontal);
+	});
+
+	test('Inserting and removing components from a container should work correctly', () => {
+		// Set up the mock proxy to save the component that gets initialized so that it can be verified
+		let rootComponent: IComponentShape;
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).callback((handle, componentShape) => rootComponent = componentShape);
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$removeFromContainer(It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+
+		// Set up the form with a top level component and a group
+		let listBox = modelView.modelBuilder.listBox().component();
+		let inputBox = modelView.modelBuilder.inputBox().component();
+		let dropDown = modelView.modelBuilder.dropDown().component();
+
+		let flex = modelView.modelBuilder.flexContainer().withItems([listBox, inputBox]).component();
+		modelView.initializeModel(flex);
+
+		assert.equal((flex as IWithItemConfig).itemConfigs.length, 2);
+		flex.insertItem(dropDown, 1);
+		assert.equal((flex as IWithItemConfig).itemConfigs.length, 3);
+		assert.equal((flex as IWithItemConfig).itemConfigs[1].toIItemConfig().componentShape.type, ModelComponentTypes.DropDown);
+		flex.removeItem(listBox);
+		assert.equal((flex as IWithItemConfig).itemConfigs.length, 2);
+		assert.equal((flex as IWithItemConfig).itemConfigs[0].toIItemConfig().componentShape.type, ModelComponentTypes.DropDown);
+	});
+
+	test('Inserting component give negative number fails', () => {
+		// Set up the mock proxy to save the component that gets initialized so that it can be verified
+		let rootComponent: IComponentShape;
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).callback((handle, componentShape) => rootComponent = componentShape);
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$removeFromContainer(It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+
+		// Set up the form with a top level component and a group
+		let listBox = modelView.modelBuilder.listBox().component();
+		let inputBox = modelView.modelBuilder.inputBox().component();
+		let dropDown = modelView.modelBuilder.dropDown().component();
+
+		let flex = modelView.modelBuilder.flexContainer().withItems([listBox, inputBox]).component();
+		modelView.initializeModel(flex);
+
+		assert.equal((flex as IWithItemConfig).itemConfigs.length, 2);
+		assert.throws(() => flex.insertItem(dropDown, -1));
+	});
+
+	test('Inserting component give wrong number fails', () => {
+		// Set up the mock proxy to save the component that gets initialized so that it can be verified
+		let rootComponent: IComponentShape;
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).callback((handle, componentShape) => rootComponent = componentShape);
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$removeFromContainer(It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+
+		// Set up the form with a top level component and a group
+		let listBox = modelView.modelBuilder.listBox().component();
+		let inputBox = modelView.modelBuilder.inputBox().component();
+		let dropDown = modelView.modelBuilder.dropDown().component();
+
+		let flex = modelView.modelBuilder.flexContainer().withItems([listBox, inputBox]).component();
+		modelView.initializeModel(flex);
+
+		assert.equal((flex as IWithItemConfig).itemConfigs.length, 2);
+		assert.throws(() => flex.insertItem(dropDown, 10));
+	});
+
+	test('Inserting component give end of the list fails', () => {
+		// Set up the mock proxy to save the component that gets initialized so that it can be verified
+		let rootComponent: IComponentShape;
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).callback((handle, componentShape) => rootComponent = componentShape);
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$removeFromContainer(It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+
+		// Set up the form with a top level component and a group
+		let listBox = modelView.modelBuilder.listBox().component();
+		let inputBox = modelView.modelBuilder.inputBox().component();
+		let dropDown = modelView.modelBuilder.dropDown().component();
+
+		let flex = modelView.modelBuilder.flexContainer().withItems([listBox, inputBox]).component();
+		modelView.initializeModel(flex);
+
+		assert.equal((flex as IWithItemConfig).itemConfigs.length, 2);
+		assert.throws(() => flex.insertItem(dropDown, 2));
+	});
+
+	test('Removing a component that does not exist does not fail', () => {
+		// Set up the mock proxy to save the component that gets initialized so that it can be verified
+		let rootComponent: IComponentShape;
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).callback((handle, componentShape) => rootComponent = componentShape);
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+
+		// Set up the form with a top level component and a group
+		let listBox = modelView.modelBuilder.listBox().component();
+		let inputBox = modelView.modelBuilder.inputBox().component();
+		let dropDown = modelView.modelBuilder.dropDown().component();
+
+		let flex = modelView.modelBuilder.flexContainer().withItems([listBox, inputBox]).component();
+		modelView.initializeModel(flex);
+
+		let result = flex.removeItem(dropDown);
+		assert.equal(result, false);
+		assert.equal((flex as IWithItemConfig).itemConfigs.length, 2);
+		assert.equal((flex as IWithItemConfig).itemConfigs[0].toIItemConfig().componentShape.type, ModelComponentTypes.ListBox);
+	});
+
+
+	test('Inserting and removing component in a form should work correctly', () => {
+		// Set up the mock proxy to save the component that gets initialized so that it can be verified
+		let rootComponent: IComponentShape;
+		mockProxy.setup(x => x.$initializeModel(It.isAny(), It.isAny())).callback((handle, componentShape) => rootComponent = componentShape);
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), undefined)).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$addToContainer(It.isAny(), It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+		mockProxy.setup(x => x.$removeFromContainer(It.isAny(), It.isAny(), It.isAny())).returns(() => Promise.resolve());
+
+		// Set up the form with a top level component and a group
+		let listBox = modelView.modelBuilder.listBox().component();
+		let inputBox = modelView.modelBuilder.inputBox().component();
+		let dropDown = modelView.modelBuilder.dropDown().component();
+		let checkBox = modelView.modelBuilder.checkBox().component();
+
+		let groupItems: azdata.FormComponentGroup = {
+			title: 'Group',
+			components: [{
+				title: 'Drop Down',
+				component: dropDown
+			}, {
+				title: 'Check Box',
+				component: checkBox
+			}]
+		};
+		let listBoxFormItem: azdata.FormComponent = {
+			title: 'List Box',
+			component: listBox
+		};
+		let inputBoxFormItem: azdata.FormComponent = {
+			title: 'Input Box',
+			component: inputBox
+		};
+
+		let formBuilder = modelView.modelBuilder.formContainer();
+		formBuilder.addFormItem(listBoxFormItem);
+		let form = formBuilder.component();
+		modelView.initializeModel(formBuilder.component());
+
+		assert.equal((form as IWithItemConfig).itemConfigs.length, 1);
+		formBuilder.insertFormItem(inputBoxFormItem, 0);
+		assert.equal((form as IWithItemConfig).itemConfigs.length, 2);
+		assert.equal((form as IWithItemConfig).itemConfigs[0].toIItemConfig().componentShape.type, ModelComponentTypes.InputBox);
+		formBuilder.insertFormItem(groupItems, 0);
+		assert.equal((form as IWithItemConfig).itemConfigs.length, 5);
+		formBuilder.removeFormItem(listBoxFormItem);
+		assert.equal((form as IWithItemConfig).itemConfigs.length, 4);
+		formBuilder.removeFormItem(groupItems);
+		assert.equal((form as IWithItemConfig).itemConfigs.length, 1);
+		formBuilder.addFormItem(listBoxFormItem);
+		assert.equal((form as IWithItemConfig).itemConfigs.length, 2);
+		assert.equal((form as IWithItemConfig).itemConfigs[1].toIItemConfig().componentShape.type, ModelComponentTypes.ListBox);
+	});
+});

src/sql/workbench/test/electron-browser/api/extHostModelViewDialog.test.ts

@@ -0,0 +1,329 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import * as assert from 'assert';
+import { Mock, It, Times } from 'typemoq';
+import { ExtHostModelViewDialog } from 'sql/workbench/api/common/extHostModelViewDialog';
+import { MainThreadModelViewDialogShape, ExtHostModelViewShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
+import { MessageLevel } from 'sql/workbench/api/common/sqlExtHostTypes';
+
+
+suite('ExtHostModelViewDialog Tests', () => {
+	let extHostModelViewDialog: ExtHostModelViewDialog;
+	let mockProxy: Mock<MainThreadModelViewDialogShape>;
+	let extHostModelView: Mock<ExtHostModelViewShape>;
+
+	setup(() => {
+		mockProxy = Mock.ofInstance(<MainThreadModelViewDialogShape>{
+			$openDialog: handle => undefined,
+			$closeDialog: handle => undefined,
+			$setDialogDetails: (handle, details) => undefined,
+			$setTabDetails: (handle, details) => undefined,
+			$setButtonDetails: (handle, details) => undefined,
+			$openWizard: handle => undefined,
+			$closeWizard: handle => undefined,
+			$setWizardPageDetails: (handle, details) => undefined,
+			$setWizardDetails: (handle, details) => undefined
+		});
+		let mainContext = <IMainContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+
+		extHostModelView = Mock.ofInstance(<ExtHostModelViewShape>{
+			$registerProvider: (widget, handler, extensionLocation) => undefined
+		});
+		extHostModelViewDialog = new ExtHostModelViewDialog(mainContext, extHostModelView.object, undefined);
+	});
+
+	test('Creating a dialog returns a dialog with initialized ok and cancel buttons and the given title', () => {
+		let title = 'dialog_title';
+		let dialog = extHostModelViewDialog.createDialog(title);
+
+		assert.equal(dialog.title, title);
+		assert.equal(dialog.okButton.enabled, true);
+		assert.equal(dialog.cancelButton.enabled, true);
+	});
+
+	test('Creating a tab returns a tab with the given title', () => {
+		let title = 'tab_title';
+		let tab = extHostModelViewDialog.createTab(title);
+
+		assert.equal(tab.title, title);
+	});
+
+	test('Creating a button returns an enabled button with the given label', () => {
+		let label = 'button_label';
+		let button = extHostModelViewDialog.createButton(label);
+
+		assert.equal(button.label, label);
+		assert.equal(button.enabled, true);
+	});
+
+	test('Opening a dialog updates its tabs and buttons on the main thread', () => {
+		mockProxy.setup(x => x.$openDialog(It.isAny()));
+		mockProxy.setup(x => x.$setDialogDetails(It.isAny(), It.isAny()));
+		mockProxy.setup(x => x.$setTabDetails(It.isAny(), It.isAny()));
+		mockProxy.setup(x => x.$setButtonDetails(It.isAny(), It.isAny()));
+
+		// Create a dialog with 2 tabs and 2 custom buttons
+		let dialogTitle = 'dialog_title';
+		let dialog = extHostModelViewDialog.createDialog(dialogTitle);
+		let tab1Title = 'tab_1';
+		let tab1 = extHostModelViewDialog.createTab(tab1Title);
+		let tab2Title = 'tab_2';
+		let tab2 = extHostModelViewDialog.createTab(tab2Title);
+		dialog.content = [tab1, tab2];
+		let button1Label = 'button_1';
+		let button1 = extHostModelViewDialog.createButton(button1Label);
+		button1.enabled = false;
+		let button2Label = 'button_2';
+		let button2 = extHostModelViewDialog.createButton(button2Label);
+		dialog.customButtons = [button1, button2];
+
+		// Open the dialog and verify that the correct main thread methods were called
+		extHostModelViewDialog.openDialog(dialog);
+		mockProxy.verify(x => x.$setButtonDetails(It.isAny(), It.is(details => {
+			return details.enabled === false && details.label === button1Label;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setButtonDetails(It.isAny(), It.is(details => {
+			return details.enabled === true && details.label === button2Label;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setTabDetails(It.isAny(), It.is(details => {
+			return details.title === tab1Title;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setTabDetails(It.isAny(), It.is(details => {
+			return details.title === tab2Title;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setDialogDetails(It.isAny(), It.is(details => {
+			return details.title === dialogTitle && details.content.length === 2 && details.customButtons.length === 2;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$openDialog(It.isAny()), Times.once());
+	});
+
+	test('Button clicks are forwarded to the correct button', () => {
+		// Set up the proxy to record button handles
+		let handles = [];
+		mockProxy.setup(x => x.$setButtonDetails(It.isAny(), It.isAny())).callback((handle, details) => handles.push(handle));
+
+		// Set up the buttons to record click events
+		let label1 = 'button_1';
+		let label2 = 'button_2';
+		let button1 = extHostModelViewDialog.createButton(label1);
+		let button2 = extHostModelViewDialog.createButton(label2);
+		let clickEvents = [];
+		button1.onClick(() => clickEvents.push(1));
+		button2.onClick(() => clickEvents.push(2));
+		extHostModelViewDialog.updateButton(button1);
+		extHostModelViewDialog.updateButton(button2);
+
+		// If the main thread sends some notifications that the buttons have been clicked
+		extHostModelViewDialog.$onButtonClick(handles[0]);
+		extHostModelViewDialog.$onButtonClick(handles[1]);
+		extHostModelViewDialog.$onButtonClick(handles[1]);
+		extHostModelViewDialog.$onButtonClick(handles[0]);
+
+		// Then the clicks should have been handled by the expected handlers
+		assert.deepEqual(clickEvents, [1, 2, 2, 1]);
+	});
+
+	test('Creating a wizard returns a wizard with initialized buttons and the given title', () => {
+		let title = 'wizard_title';
+		let wizard = extHostModelViewDialog.createWizard(title);
+
+		assert.equal(wizard.title, title);
+		assert.equal(wizard.doneButton.enabled, true);
+		assert.equal(wizard.cancelButton.enabled, true);
+		assert.equal(wizard.nextButton.enabled, true);
+		assert.equal(wizard.backButton.enabled, true);
+		assert.deepEqual(wizard.pages, []);
+	});
+
+	test('Opening a wizard updates its pages and buttons on the main thread', () => {
+		mockProxy.setup(x => x.$openWizard(It.isAny()));
+		mockProxy.setup(x => x.$setWizardDetails(It.isAny(), It.isAny()));
+		mockProxy.setup(x => x.$setWizardPageDetails(It.isAny(), It.isAny()));
+		mockProxy.setup(x => x.$setButtonDetails(It.isAny(), It.isAny()));
+
+		// Create a wizard with 2 pages and 2 custom buttons
+		let wizardTitle = 'wizard_title';
+		let wizard = extHostModelViewDialog.createWizard(wizardTitle);
+		let page1Title = 'page_1';
+		let page1 = extHostModelViewDialog.createWizardPage(page1Title);
+		let page2Title = 'page_2';
+		let page2 = extHostModelViewDialog.createWizardPage(page2Title);
+		wizard.pages = [page1, page2];
+		let button1Label = 'button_1';
+		let button1 = extHostModelViewDialog.createButton(button1Label);
+		button1.enabled = false;
+		let button2Label = 'button_2';
+		let button2 = extHostModelViewDialog.createButton(button2Label);
+		wizard.customButtons = [button1, button2];
+
+		// Open the wizard and verify that the correct main thread methods were called
+		extHostModelViewDialog.openWizard(wizard);
+		mockProxy.verify(x => x.$setButtonDetails(It.isAny(), It.is(details => {
+			return details.enabled === false && details.label === button1Label;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setButtonDetails(It.isAny(), It.is(details => {
+			return details.enabled === true && details.label === button2Label;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setWizardPageDetails(It.isAny(), It.is(details => {
+			return details.title === page1Title;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setWizardPageDetails(It.isAny(), It.is(details => {
+			return details.title === page2Title;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$setWizardDetails(It.isAny(), It.is(details => {
+			return details.title === wizardTitle && details.pages.length === 2 && details.customButtons.length === 2 &&
+				details.displayPageTitles === true;
+		})), Times.atLeastOnce());
+		mockProxy.verify(x => x.$openWizard(It.isAny()), Times.once());
+	});
+
+	test('Wizard page changed events are handled correctly', () => {
+		// Set up the main thread mock to record the handle assigned to the wizard
+		let wizardHandle: number;
+		mockProxy.setup(x => x.$setWizardDetails(It.isAny(), It.isAny())).callback((handle, details) => wizardHandle = handle);
+
+		// Set up the wizard with 2 pages
+		let wizard = extHostModelViewDialog.createWizard('test_wizard');
+		let page1 = extHostModelViewDialog.createWizardPage('page_1');
+		let page2 = extHostModelViewDialog.createWizardPage('page_2');
+		wizard.pages = [page1, page2];
+		extHostModelViewDialog.updateWizard(wizard);
+
+		// Record page changed events
+		let actualPageChangeInfo = [];
+		wizard.onPageChanged(pageChangeInfo => {
+			actualPageChangeInfo.push(pageChangeInfo);
+		});
+
+		// Call the page changed event and verify that it was handled
+		let expectedPageChangeInfo = {
+			lastPage: 0,
+			newPage: 1
+		};
+		extHostModelViewDialog.$onWizardPageChanged(wizardHandle, expectedPageChangeInfo);
+		assert.equal(actualPageChangeInfo.length, 1);
+		assert.equal(actualPageChangeInfo[0], expectedPageChangeInfo);
+		assert.equal(wizard.currentPage, expectedPageChangeInfo.newPage);
+	});
+
+	test('Validity changed events are handled correctly', () => {
+		// Set up the main thread mock to record handles assigned to tabs
+		let tabHandles = [];
+		mockProxy.setup(x => x.$setTabDetails(It.isAny(), It.isAny())).callback((handle, details) => tabHandles.push(handle));
+
+		// Set up the dialog with 2 tabs
+		let dialog = extHostModelViewDialog.createDialog('test_dialog');
+		let tab1 = extHostModelViewDialog.createTab('tab_1');
+		let tab2 = extHostModelViewDialog.createTab('tab_2');
+		dialog.content = [tab1, tab2];
+		extHostModelViewDialog.updateDialogContent(dialog);
+
+		// Record tab validity changed events
+		let tab1ValidityChangedEvents = [];
+		let tab2ValidityChangedEvents = [];
+		tab1.onValidityChanged(valid => tab1ValidityChangedEvents.push(valid));
+		tab2.onValidityChanged(valid => tab2ValidityChangedEvents.push(valid));
+
+		// Call the validity changed event on tab 2 and verify that it was handled but tab 1 is still not valid
+		extHostModelViewDialog.$onPanelValidityChanged(tabHandles[1], false);
+		assert.equal(tab1ValidityChangedEvents.length, 0);
+		assert.equal(tab1.valid, true);
+		assert.equal(tab2ValidityChangedEvents.length, 1);
+		assert.equal(tab2ValidityChangedEvents[0], false);
+		assert.equal(tab2.valid, false);
+	});
+
+	test('Verify validity changed events update validity for all panel types', () => {
+		// Set up the main thread mock to record handles for the tab, dialog, and page
+		let tabHandle: number;
+		let dialogHandle: number;
+		let pageHandle: number;
+		mockProxy.setup(x => x.$setTabDetails(It.isAny(), It.isAny())).callback((handle, details) => tabHandle = handle);
+		mockProxy.setup(x => x.$setDialogDetails(It.isAny(), It.isAny())).callback((handle, details) => dialogHandle = handle);
+		mockProxy.setup(x => x.$setWizardPageDetails(It.isAny(), It.isAny())).callback((handle, details) => pageHandle = handle);
+
+		// Initialize a tab, dialog, and page
+		let tab = extHostModelViewDialog.createTab('tab_1');
+		extHostModelViewDialog.updateTabContent(tab);
+		let dialog = extHostModelViewDialog.createDialog('dialog_1');
+		extHostModelViewDialog.updateDialogContent(dialog);
+		let page = extHostModelViewDialog.createWizardPage('page_1');
+		extHostModelViewDialog.updateWizardPage(page);
+
+		// Call the validity changed event on each object and verify that the object's validity was updated
+		extHostModelViewDialog.$onPanelValidityChanged(tabHandle, false);
+		assert.equal(tab.valid, false);
+		extHostModelViewDialog.$onPanelValidityChanged(dialogHandle, false);
+		assert.equal(dialog.valid, false);
+		extHostModelViewDialog.$onPanelValidityChanged(pageHandle, false);
+		assert.equal(page.valid, false);
+	});
+
+	test('Main thread can execute wizard navigation validation', () => {
+		// Set up the main thread mock to record the wizard handle
+		let wizardHandle: number;
+		mockProxy.setup(x => x.$setWizardDetails(It.isAny(), It.isAny())).callback((handle, details) => wizardHandle = handle);
+
+		// Create the wizard and add a validation that records that it has been called
+		let wizard = extHostModelViewDialog.createWizard('wizard_1');
+		extHostModelViewDialog.updateWizard(wizard);
+		let validationInfo: azdata.window.WizardPageChangeInfo;
+		wizard.registerNavigationValidator(info => {
+			validationInfo = info;
+			return true;
+		});
+
+		// If I call the validation from the main thread then it should run and record the correct page change info
+		let lastPage = 0;
+		let newPage = 1;
+		extHostModelViewDialog.$validateNavigation(wizardHandle, {
+			lastPage: lastPage,
+			newPage: newPage
+		});
+		assert.notEqual(validationInfo, undefined);
+		assert.equal(validationInfo.lastPage, lastPage);
+		assert.equal(validationInfo.newPage, newPage);
+	});
+
+	test('Changing the wizard message sends the new message to the main thread', () => {
+		// Set up the main thread mock to record the call
+		mockProxy.setup(x => x.$setWizardDetails(It.isAny(), It.isAny()));
+		let wizard = extHostModelViewDialog.createWizard('wizard_1');
+
+		// If I update the wizard's message
+		let newMessage = {
+			level: MessageLevel.Error,
+			text: 'test message'
+		};
+		wizard.message = newMessage;
+
+		// Then the main thread gets notified of the new details
+		mockProxy.verify(x => x.$setWizardDetails(It.isAny(), It.is(x => x.message === newMessage)), Times.once());
+	});
+
+	test('Main thread can execute dialog close validation', () => {
+		// Set up the main thread mock to record the dialog handle
+		let dialogHandle: number;
+		mockProxy.setup(x => x.$setDialogDetails(It.isAny(), It.isAny())).callback((handle, details) => dialogHandle = handle);
+
+		// Create the dialog and add a validation that records that it has been called
+		let dialog = extHostModelViewDialog.createDialog('dialog_1');
+		extHostModelViewDialog.updateDialogContent(dialog);
+		let callCount = 0;
+		dialog.registerCloseValidator(() => {
+			callCount++;
+			return true;
+		});
+
+		// If I call the validation from the main thread then it should run
+		extHostModelViewDialog.$validateDialogClose(dialogHandle);
+		assert.equal(callCount, 1);
+	});
+});

src/sql/workbench/test/electron-browser/api/extHostObjectExplorer.test.ts

@@ -0,0 +1,111 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import * as assert from 'assert';
+import * as TypeMoq from 'typemoq';
+
+import { MainThreadObjectExplorerShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { ExtHostObjectExplorerNode } from 'sql/workbench/api/common/extHostObjectExplorer';
+
+const nodes: { [nodeName: string]: azdata.NodeInfo } =
+{
+	'Server1': {
+		nodePath: 'MyServer',
+		nodeStatus: '',
+		nodeSubType: '',
+		nodeType: 'Server',
+		isLeaf: false,
+		label: 'MyServer',
+		metadata: undefined,
+		errorMessage: ''
+	},
+	'DatabasesFolder': {
+		nodePath: 'MyServer/Databases',
+		nodeStatus: '',
+		nodeSubType: '',
+		nodeType: 'Folder',
+		isLeaf: false,
+		label: 'Databases',
+		metadata: undefined,
+		errorMessage: ''
+	},
+	'Database1': {
+		nodePath: 'MyServer/Databases/MyDatabase',
+		nodeStatus: '',
+		nodeSubType: '',
+		nodeType: 'Database',
+		isLeaf: false,
+		label: 'MyDatabase',
+		metadata: undefined,
+		errorMessage: ''
+	},
+	'Database2': {
+		nodePath: 'MyServer/Databases/My/TrickyDatabase',
+		nodeStatus: '',
+		nodeSubType: '',
+		nodeType: 'Database',
+		isLeaf: false,
+		label: 'My/TrickyDatabase',
+		metadata: undefined,
+		errorMessage: ''
+	},
+	'TablesFolder': {
+		nodePath: 'MyServer/Databases/My/TrickyDatabase/Tables',
+		nodeStatus: '',
+		nodeSubType: '',
+		nodeType: 'Folder',
+		isLeaf: false,
+		label: 'Tables',
+		metadata: undefined,
+		errorMessage: ''
+	}
+};
+
+suite('ExtHostObjectExplorer Tests', () => {
+	let mockProxy: TypeMoq.Mock<MainThreadObjectExplorerShape>;
+	suiteSetup(() => {
+		mockProxy = TypeMoq.Mock.ofInstance(<MainThreadObjectExplorerShape>{
+			$getNode: (connectionId: string, nodePath?: string): Thenable<azdata.NodeInfo> => undefined
+		});
+
+		mockProxy.setup(p =>
+			p.$getNode(TypeMoq.It.isAny(), TypeMoq.It.isAny()))
+			.returns((connectionId, nodePath) => {
+				return Promise.resolve<azdata.NodeInfo>(nodes[Object.keys(nodes).find(key =>
+					nodes[key].nodePath === nodePath)]);
+			});
+	});
+
+	suite('ExtHostObjectExplorerNode', () => {
+		let extHostObjectExplorerNode: ExtHostObjectExplorerNode;
+		suite('getParent', () => {
+			test('Should return undefined if no parent', async () => {
+				extHostObjectExplorerNode = new ExtHostObjectExplorerNode(nodes['Server1'], 'connectionId', mockProxy.object);
+				assert.equal(await extHostObjectExplorerNode.getParent(), undefined);
+			});
+
+			test('should return root with direct descendent of root', async () => {
+				extHostObjectExplorerNode = new ExtHostObjectExplorerNode(nodes['DatabasesFolder'], 'connectionId', mockProxy.object);
+				assert.equal((await extHostObjectExplorerNode.getParent()).nodePath, nodes['Server1'].nodePath);
+			});
+
+			test('should return correct parent with further descendent of root', async () => {
+				extHostObjectExplorerNode = new ExtHostObjectExplorerNode(nodes['Database1'], 'connectionId', mockProxy.object);
+				assert.equal((await extHostObjectExplorerNode.getParent()).nodePath, nodes['DatabasesFolder'].nodePath);
+			});
+
+			test('should return correct parent with node having / in its name', async () => {
+				extHostObjectExplorerNode = new ExtHostObjectExplorerNode(nodes['Database2'], 'connectionId', mockProxy.object);
+				assert.equal((await extHostObjectExplorerNode.getParent()).nodePath, nodes['DatabasesFolder'].nodePath);
+			});
+
+			test('should return correct parent with parent node having / in its name', async () => {
+				extHostObjectExplorerNode = new ExtHostObjectExplorerNode(nodes['TablesFolder'], 'connectionId', mockProxy.object);
+				assert.equal((await extHostObjectExplorerNode.getParent()).nodePath, nodes['Database2'].nodePath);
+			});
+		});
+	});
+});

src/sql/workbench/test/electron-browser/api/exthostNotebook.test.ts

@@ -0,0 +1,147 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  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 assert from 'assert';
+import * as TypeMoq from 'typemoq';
+
+import { URI } from 'vs/base/common/uri';
+import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
+
+import { ExtHostNotebook } from 'sql/workbench/api/common/extHostNotebook';
+import { MainThreadNotebookShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import * as testUtils from '../../../../../sqltest/utils/testUtils';
+import { INotebookManagerDetails } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { mssqlProviderName } from 'sql/platform/connection/common/constants';
+
+suite('ExtHostNotebook Tests', () => {
+
+	let extHostNotebook: ExtHostNotebook;
+	let mockProxy: TypeMoq.Mock<MainThreadNotebookShape>;
+	let notebookUri: URI;
+	let notebookProviderMock: TypeMoq.Mock<NotebookProviderStub>;
+	setup(() => {
+		mockProxy = TypeMoq.Mock.ofInstance(<MainThreadNotebookShape>{
+			$registerNotebookProvider: (providerId, handle) => undefined,
+			$unregisterNotebookProvider: (handle) => undefined,
+			dispose: () => undefined
+		});
+		let mainContext = <IMainContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+		extHostNotebook = new ExtHostNotebook(mainContext);
+		notebookUri = URI.parse('file:/user/default/my.ipynb');
+		notebookProviderMock = TypeMoq.Mock.ofType(NotebookProviderStub, TypeMoq.MockBehavior.Loose);
+		notebookProviderMock.callBase = true;
+	});
+
+	suite('getNotebookManager', () => {
+		test('Should throw if no matching provider is defined', async () => {
+			await testUtils.assertThrowsAsync(() => extHostNotebook.$getNotebookManager(-1, notebookUri));
+		});
+		suite('with provider', () => {
+			let providerHandle: number = -1;
+
+			setup(() => {
+				mockProxy.setup(p =>
+					p.$registerNotebookProvider(TypeMoq.It.isValue(notebookProviderMock.object.providerId), TypeMoq.It.isAnyNumber()))
+					.returns((providerId, handle) => {
+						providerHandle = handle;
+						return undefined;
+					});
+
+				// Register the provider so we can test behavior with this present
+				extHostNotebook.registerNotebookProvider(notebookProviderMock.object);
+			});
+
+			test('Should return a notebook manager with correct info on content and server manager existence', async () => {
+				// Given the provider returns a manager with no
+				let expectedManager = new NotebookManagerStub();
+				notebookProviderMock.setup(p => p.getNotebookManager(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedManager));
+
+				// When I call through using the handle provided during registration
+				let managerDetails: INotebookManagerDetails = await extHostNotebook.$getNotebookManager(providerHandle, notebookUri);
+
+				// Then I expect the same manager to be returned
+				assert.ok(managerDetails.hasContentManager === false, 'Expect no content manager defined');
+				assert.ok(managerDetails.hasServerManager === false, 'Expect no server manager defined');
+				assert.ok(managerDetails.handle > 0, 'Expect a valid handle defined');
+			});
+
+			test('Should have a unique handle for each notebook URI', async () => {
+				// Given the we request 2 URIs
+				let expectedManager = new NotebookManagerStub();
+				notebookProviderMock.setup(p => p.getNotebookManager(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedManager));
+
+				// When I call through using the handle provided during registration
+				let originalManagerDetails = await extHostNotebook.$getNotebookManager(providerHandle, notebookUri);
+				let differentDetails = await extHostNotebook.$getNotebookManager(providerHandle, URI.parse('file://other/file.ipynb'));
+				let sameDetails = await extHostNotebook.$getNotebookManager(providerHandle, notebookUri);
+
+				// Then I expect the 2 different handles in the managers returned.
+				// This is because we can't easily track identity of the managers, so just track which one is assigned to
+				// a notebook by the handle ID
+				assert.notEqual(originalManagerDetails.handle, differentDetails.handle, 'Should have unique handle for each manager');
+				assert.equal(originalManagerDetails.handle, sameDetails.handle, 'Should have same handle when same URI is passed in');
+
+			});
+		});
+	});
+
+	suite('registerNotebookProvider', () => {
+		let savedHandle: number = -1;
+		setup(() => {
+			mockProxy.setup(p =>
+				p.$registerNotebookProvider(TypeMoq.It.isValue(notebookProviderMock.object.providerId), TypeMoq.It.isAnyNumber()))
+				.returns((providerId, handle) => {
+					savedHandle = handle;
+					return undefined;
+				});
+		});
+
+		test('Should register with a new handle to the proxy', () => {
+			extHostNotebook.registerNotebookProvider(notebookProviderMock.object);
+			mockProxy.verify(p =>
+				p.$registerNotebookProvider(TypeMoq.It.isValue(notebookProviderMock.object.providerId),
+					TypeMoq.It.isAnyNumber()), TypeMoq.Times.once());
+			// It shouldn't unregister until requested
+			mockProxy.verify(p => p.$unregisterNotebookProvider(TypeMoq.It.isValue(savedHandle)), TypeMoq.Times.never());
+
+		});
+
+		test('Should not call unregister on disposing', () => {
+			let disposable = extHostNotebook.registerNotebookProvider(notebookProviderMock.object);
+			disposable.dispose();
+			mockProxy.verify(p => p.$unregisterNotebookProvider(TypeMoq.It.isValue(savedHandle)), TypeMoq.Times.never());
+		});
+	});
+});
+
+class NotebookProviderStub implements azdata.nb.NotebookProvider {
+	providerId: string = 'TestProvider';
+	standardKernels: azdata.nb.IStandardKernel[] = [{ name: 'fakeKernel', displayName: 'fakeKernel', connectionProviderIds: [mssqlProviderName] }];
+
+	getNotebookManager(notebookUri: vscode.Uri): Thenable<azdata.nb.NotebookManager> {
+		throw new Error('Method not implemented.');
+	}
+	handleNotebookClosed(notebookUri: vscode.Uri): void {
+		throw new Error('Method not implemented.');
+	}
+}
+
+class NotebookManagerStub implements azdata.nb.NotebookManager {
+	get contentManager(): azdata.nb.ContentManager {
+		return undefined;
+	}
+
+	get sessionManager(): azdata.nb.SessionManager {
+		return undefined;
+	}
+
+	get serverManager(): azdata.nb.ServerManager {
+		return undefined;
+	}
+}

src/sql/workbench/test/electron-browser/api/mainThreadBackgroundTaskManagement.test.ts

@@ -0,0 +1,94 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import { Mock, It, Times } from 'typemoq';
+import { MainThreadBackgroundTaskManagement, TaskStatus } from 'sql/workbench/api/browser/mainThreadBackgroundTaskManagement';
+import { ITaskService } from 'sql/platform/tasks/common/tasksService';
+import { IExtHostContext } from 'vs/workbench/api/common/extHost.protocol';
+import { TaskNode } from 'sql/platform/tasks/common/tasksNode';
+import { Emitter } from 'vs/base/common/event';
+import { ExtHostBackgroundTaskManagementShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+
+suite('MainThreadBackgroundTaskManagement Tests', () => {
+	let mainThreadBackgroundTaskManagement: MainThreadBackgroundTaskManagement;
+	let mockProxy: Mock<ExtHostBackgroundTaskManagementShape>;
+	let taskService: Mock<ITaskService>;
+	let nothing: void;
+	let operationId = 'operation is';
+	let onTaskComplete = new Emitter<TaskNode>();
+	setup(() => {
+		mockProxy = Mock.ofInstance(<ExtHostBackgroundTaskManagementShape>{
+			$onTaskRegistered: (operationId: string) => nothing,
+			$onTaskCanceled: (operationId: string) => nothing,
+			$registerTask: (operationInfo: azdata.BackgroundOperationInfo) => nothing,
+			$removeTask: (operationId: string) => nothing,
+		});
+		taskService = Mock.ofInstance(<ITaskService>{
+			_serviceBrand: undefined,
+			onTaskComplete: undefined,
+			onAddNewTask: undefined,
+			handleNewTask: undefined,
+			handleTaskComplete: undefined,
+			getAllTasks: undefined,
+			getNumberOfInProgressTasks: undefined,
+			onNewTaskCreated: undefined,
+			createNewTask: (taskInfo: azdata.TaskInfo) => nothing,
+			updateTask: (taskProgressInfo: azdata.TaskProgressInfo) => nothing,
+			onTaskStatusChanged: undefined,
+			cancelTask: undefined,
+			registerProvider: undefined
+		});
+		let mainContext = <IExtHostContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+
+		taskService.setup(x => x.onTaskComplete).returns(() => onTaskComplete.event);
+
+		mainThreadBackgroundTaskManagement = new MainThreadBackgroundTaskManagement(mainContext, taskService.object);
+	});
+
+	test('RegisterTask should successfully create background task', () => {
+		let taskInfo: azdata.TaskInfo = {
+			taskId: operationId,
+			databaseName: undefined,
+			description: undefined,
+			isCancelable: true,
+			name: 'task name',
+			providerName: undefined,
+			serverName: undefined,
+			status: TaskStatus.NotStarted,
+			taskExecutionMode: 0
+		};
+		mainThreadBackgroundTaskManagement.$registerTask(taskInfo);
+		taskService.verify(x => x.createNewTask(It.is(t => t.status === TaskStatus.NotStarted)), Times.once());
+		mockProxy.verify(x => x.$onTaskRegistered(operationId), Times.once());
+	});
+
+	test('UpdateTask should successfully update the background task status', () => {
+		let taskInfo: azdata.TaskProgressInfo = {
+			taskId: operationId,
+			status: TaskStatus.InProgress,
+			message: undefined,
+		};
+		mainThreadBackgroundTaskManagement.$updateTask(taskInfo);
+		taskService.verify(x => x.updateTask(It.is(t => t.status === TaskStatus.InProgress)), Times.once());
+	});
+
+	test('Canceling the task should notify the proxy', () => {
+		let taskInfo: azdata.TaskProgressInfo = {
+			taskId: operationId,
+			status: TaskStatus.InProgress,
+			message: undefined,
+		};
+		let taskNode = new TaskNode('', '', '', operationId, undefined);
+		taskNode.status = TaskStatus.Canceling;
+
+		onTaskComplete.fire(taskNode);
+		mainThreadBackgroundTaskManagement.$updateTask(taskInfo);
+		mockProxy.verify(x => x.$onTaskCanceled(It.is(t => t === operationId)), Times.once());
+	});
+
+});

src/sql/workbench/test/electron-browser/api/mainThreadModelViewDialog.test.ts

@@ -0,0 +1,371 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import { Mock, It, Times } from 'typemoq';
+import { MainThreadModelViewDialog } from 'sql/workbench/api/electron-browser/mainThreadModelViewDialog';
+import { IExtHostContext } from 'vs/workbench/api/common/extHost.protocol';
+import { IModelViewButtonDetails, IModelViewTabDetails, IModelViewDialogDetails, IModelViewWizardPageDetails, IModelViewWizardDetails, DialogMessage, MessageLevel } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { CustomDialogService } from 'sql/platform/dialog/browser/customDialogService';
+import { Dialog, DialogTab, Wizard } from 'sql/platform/dialog/common/dialogTypes';
+import { ExtHostModelViewDialogShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { Emitter } from 'vs/base/common/event';
+
+
+suite('MainThreadModelViewDialog Tests', () => {
+	let mainThreadModelViewDialog: MainThreadModelViewDialog;
+	let mockExtHostModelViewDialog: Mock<ExtHostModelViewDialogShape>;
+	let mockDialogService: Mock<CustomDialogService>;
+	let openedDialog: Dialog;
+	let openedWizard: Wizard;
+
+	// Dialog details
+	let button1Details: IModelViewButtonDetails;
+	let button2Details: IModelViewButtonDetails;
+	let okButtonDetails: IModelViewButtonDetails;
+	let cancelButtonDetails: IModelViewButtonDetails;
+	let tab1Details: IModelViewTabDetails;
+	let tab2Details: IModelViewTabDetails;
+	let dialogDetails: IModelViewDialogDetails;
+	let button1Handle = 1;
+	let button2Handle = 2;
+	let okButtonHandle = 3;
+	let cancelButtonHandle = 4;
+	let tab1Handle = 5;
+	let tab2Handle = 6;
+	let dialogHandle = 7;
+
+	// Wizard details
+	let nextButtonDetails: IModelViewButtonDetails;
+	let backButtonDetails: IModelViewButtonDetails;
+	let generateScriptButtonDetails: IModelViewButtonDetails;
+	let page1Details: IModelViewWizardPageDetails;
+	let page2Details: IModelViewWizardPageDetails;
+	let page3Details: IModelViewWizardPageDetails;
+	let wizardDetails: IModelViewWizardDetails;
+	let nextButtonHandle = 8;
+	let backButtonHandle = 9;
+	let generateScriptButtonHandle = 10;
+	let page1Handle = 11;
+	let page2Handle = 12;
+	let wizardHandle = 13;
+	let page3Handle = 14;
+
+	setup(() => {
+		mockExtHostModelViewDialog = Mock.ofInstance(<ExtHostModelViewDialogShape>{
+			$onButtonClick: handle => undefined,
+			$onPanelValidityChanged: (handle, valid) => undefined,
+			$onWizardPageChanged: (handle, info) => undefined,
+			$updateWizardPageInfo: (wizardHandle, pageHandles, currentPageIndex) => undefined,
+			$validateNavigation: (handle, info) => undefined,
+			$validateDialogClose: handle => undefined
+		});
+		let extHostContext = <IExtHostContext>{
+			getProxy: proxyType => mockExtHostModelViewDialog.object
+		};
+		mainThreadModelViewDialog = new MainThreadModelViewDialog(extHostContext, undefined, undefined, undefined);
+
+		// Set up the mock dialog service
+		mockDialogService = Mock.ofType(CustomDialogService, undefined);
+		openedDialog = undefined;
+		mockDialogService.setup(x => x.showDialog(It.isAny(), undefined, It.isAny())).callback((dialog) => {
+			openedDialog = dialog;
+		});
+		mockDialogService.setup(x => x.showWizard(It.isAny())).callback(wizard => {
+			openedWizard = wizard;
+			// The actual service will set the page to 0 when it opens the wizard
+			openedWizard.setCurrentPage(0);
+		});
+		(mainThreadModelViewDialog as any)._dialogService = mockDialogService.object;
+
+		// Set up the dialog details
+		button1Details = {
+			label: 'button1',
+			enabled: false,
+			hidden: false
+		};
+		button2Details = {
+			label: 'button2',
+			enabled: true,
+			hidden: false
+		};
+		okButtonDetails = {
+			label: 'ok_label',
+			enabled: true,
+			hidden: false
+		};
+		cancelButtonDetails = {
+			label: 'cancel_label',
+			enabled: true,
+			hidden: false
+		};
+		tab1Details = {
+			title: 'tab1',
+			content: 'content1'
+		};
+		tab2Details = {
+			title: 'tab2',
+			content: 'content2'
+		};
+		dialogDetails = {
+			title: 'dialog1',
+			isWide: false,
+			content: [tab1Handle, tab2Handle],
+			okButton: okButtonHandle,
+			cancelButton: cancelButtonHandle,
+			customButtons: [button1Handle, button2Handle],
+			message: undefined
+		};
+
+		// Set up the wizard details
+		nextButtonDetails = {
+			label: 'next_label',
+			enabled: true,
+			hidden: false
+		};
+		backButtonDetails = {
+			label: 'back_label',
+			enabled: true,
+			hidden: false
+		};
+		generateScriptButtonDetails = {
+			label: 'generate_script_label',
+			enabled: true,
+			hidden: false
+		};
+		page1Details = {
+			title: 'page1',
+			content: 'content1',
+			enabled: true,
+			customButtons: [],
+			description: 'description1'
+		};
+		page2Details = {
+			title: 'page2',
+			content: 'content2',
+			enabled: true,
+			customButtons: [button1Handle, button2Handle],
+			description: 'description2'
+		};
+		wizardDetails = {
+			backButton: backButtonHandle,
+			nextButton: nextButtonHandle,
+			generateScriptButton: generateScriptButtonHandle,
+			cancelButton: cancelButtonHandle,
+			doneButton: okButtonHandle,
+			currentPage: undefined,
+			title: 'wizard_title',
+			customButtons: [],
+			pages: [page1Handle, page2Handle],
+			message: undefined,
+			displayPageTitles: false
+		};
+
+		// Register the buttons, tabs, and dialog
+		mainThreadModelViewDialog.$setButtonDetails(button1Handle, button1Details);
+		mainThreadModelViewDialog.$setButtonDetails(button2Handle, button2Details);
+		mainThreadModelViewDialog.$setButtonDetails(okButtonHandle, okButtonDetails);
+		mainThreadModelViewDialog.$setButtonDetails(cancelButtonHandle, cancelButtonDetails);
+		mainThreadModelViewDialog.$setTabDetails(tab1Handle, tab1Details);
+		mainThreadModelViewDialog.$setTabDetails(tab2Handle, tab2Details);
+		mainThreadModelViewDialog.$setDialogDetails(dialogHandle, dialogDetails);
+
+		// Register the wizard and its pages and buttons
+		mainThreadModelViewDialog.$setButtonDetails(nextButtonHandle, nextButtonDetails);
+		mainThreadModelViewDialog.$setButtonDetails(backButtonHandle, backButtonDetails);
+		mainThreadModelViewDialog.$setButtonDetails(generateScriptButtonHandle, generateScriptButtonDetails);
+		mainThreadModelViewDialog.$setWizardPageDetails(page1Handle, page1Details);
+		mainThreadModelViewDialog.$setWizardPageDetails(page2Handle, page2Details);
+		mainThreadModelViewDialog.$setWizardDetails(wizardHandle, wizardDetails);
+	});
+
+	test('Creating a dialog and calling open on it causes a dialog with correct content and buttons to open', () => {
+		// If I open the dialog
+		mainThreadModelViewDialog.$openDialog(dialogHandle);
+
+		// Then the opened dialog's content and buttons match what was set
+		mockDialogService.verify(x => x.showDialog(It.isAny(), undefined, It.isAny()), Times.once());
+		assert.notEqual(openedDialog, undefined);
+		assert.equal(openedDialog.title, dialogDetails.title);
+		assert.equal(openedDialog.okButton.label, okButtonDetails.label);
+		assert.equal(openedDialog.okButton.enabled, okButtonDetails.enabled);
+		assert.equal(openedDialog.cancelButton.label, cancelButtonDetails.label);
+		assert.equal(openedDialog.cancelButton.enabled, cancelButtonDetails.enabled);
+		assert.equal(openedDialog.customButtons.length, 2);
+		assert.equal(openedDialog.customButtons[0].label, button1Details.label);
+		assert.equal(openedDialog.customButtons[0].enabled, button1Details.enabled);
+		assert.equal(openedDialog.customButtons[1].label, button2Details.label);
+		assert.equal(openedDialog.customButtons[1].enabled, button2Details.enabled);
+		assert.equal(openedDialog.content.length, 2);
+		assert.equal((openedDialog.content[0] as DialogTab).content, tab1Details.content);
+		assert.equal((openedDialog.content[0] as DialogTab).title, tab1Details.title);
+		assert.equal((openedDialog.content[1] as DialogTab).content, tab2Details.content);
+		assert.equal((openedDialog.content[1] as DialogTab).title, tab2Details.title);
+	});
+
+	test('Button presses are forwarded to the extension host', () => {
+		// Set up the mock proxy to capture button presses
+		let pressedHandles = [];
+		mockExtHostModelViewDialog.setup(x => x.$onButtonClick(It.isAny())).callback(handle => pressedHandles.push(handle));
+
+		// Open the dialog so that its buttons can be accessed
+		mainThreadModelViewDialog.$openDialog(dialogHandle);
+
+		// Set up click emitters for each button
+		let okEmitter = new Emitter<void>();
+		let cancelEmitter = new Emitter<void>();
+		let button1Emitter = new Emitter<void>();
+		let button2Emitter = new Emitter<void>();
+		openedDialog.okButton.registerClickEvent(okEmitter.event);
+		openedDialog.cancelButton.registerClickEvent(cancelEmitter.event);
+		openedDialog.customButtons[0].registerClickEvent(button1Emitter.event);
+		openedDialog.customButtons[1].registerClickEvent(button2Emitter.event);
+
+		// Click the buttons
+		button1Emitter.fire();
+		button2Emitter.fire();
+		okEmitter.fire();
+		cancelEmitter.fire();
+		button2Emitter.fire();
+		cancelEmitter.fire();
+		button1Emitter.fire();
+		okEmitter.fire();
+
+		// Verify that the correct button click notifications were sent to the proxy
+		assert.deepEqual(pressedHandles, [button1Handle, button2Handle, okButtonHandle, cancelButtonHandle, button2Handle, cancelButtonHandle, button1Handle, okButtonHandle]);
+	});
+
+	test('Creating a wizard and calling open on it causes a wizard with correct pages and buttons to open', () => {
+		// If I open the wizard
+		mainThreadModelViewDialog.$openWizard(wizardHandle);
+
+		// Then the opened wizard's content and buttons match what was set
+		mockDialogService.verify(x => x.showWizard(It.isAny()), Times.once());
+		assert.notEqual(openedWizard, undefined);
+		assert.equal(openedWizard.title, wizardDetails.title);
+		assert.equal(openedWizard.doneButton.label, okButtonDetails.label);
+		assert.equal(openedWizard.doneButton.enabled, okButtonDetails.enabled);
+		assert.equal(openedWizard.cancelButton.label, cancelButtonDetails.label);
+		assert.equal(openedWizard.cancelButton.enabled, cancelButtonDetails.enabled);
+		assert.equal(openedWizard.customButtons.length, 0);
+		assert.equal(openedWizard.pages.length, 2);
+		assert.equal(openedWizard.currentPage, 0);
+		assert.equal(openedWizard.displayPageTitles, wizardDetails.displayPageTitles);
+		let page1 = openedWizard.pages[0];
+		assert.equal(page1.title, page1Details.title);
+		assert.equal(page1.content, page1Details.content);
+		assert.equal(page1.enabled, page1Details.enabled);
+		assert.equal(page1.valid, true);
+		assert.equal(page1.customButtons.length, 0);
+		assert.equal(page1.description, page1Details.description);
+		let page2 = openedWizard.pages[1];
+		assert.equal(page2.title, page2Details.title);
+		assert.equal(page2.content, page2Details.content);
+		assert.equal(page2.enabled, page2Details.enabled);
+		assert.equal(page2.valid, true);
+		assert.equal(page2.customButtons.length, 2);
+		assert.equal(page2.description, page2Details.description);
+	});
+
+	test('The extension host gets notified when wizard page change events occur', () => {
+		mockExtHostModelViewDialog.setup(x => x.$onWizardPageChanged(It.isAny(), It.isAny()));
+
+		// If I open the wizard and change the page to index 1
+		mainThreadModelViewDialog.$openWizard(wizardHandle);
+		openedWizard.setCurrentPage(1);
+
+		// Then a page changed event gets sent to the extension host
+		mockExtHostModelViewDialog.verify(x => x.$onWizardPageChanged(It.is(handle => handle === wizardHandle),
+			It.is(pageChangeInfo => pageChangeInfo.lastPage === 0 && pageChangeInfo.newPage === 1)), Times.once());
+	});
+
+	test('Validity changed events are forwarded to the extension host', () => {
+		mockExtHostModelViewDialog.setup(x => x.$onPanelValidityChanged(It.isAny(), It.isAny()));
+
+		// If I open the dialog and set its validity and its 2nd tab's validity to false
+		mainThreadModelViewDialog.$openDialog(dialogHandle);
+		(openedDialog.content[1] as DialogTab).notifyValidityChanged(false);
+		openedDialog.notifyValidityChanged(false);
+
+		// Then a validity changed event gets sent to the extension host for the tab and the dialog
+		mockExtHostModelViewDialog.verify(x => x.$onPanelValidityChanged(It.is(handle => handle === dialogHandle), It.is(valid => valid === false)), Times.once());
+		mockExtHostModelViewDialog.verify(x => x.$onPanelValidityChanged(It.is(handle => handle === tab2Handle), It.is(valid => valid === false)), Times.once());
+	});
+
+	test('addWizardPage method inserts pages at the correct spot and notifies the extension host', () => {
+		mockExtHostModelViewDialog.setup(x => x.$updateWizardPageInfo(It.isAny(), It.isAny(), It.isAny()));
+		page3Details = {
+			title: 'page_3',
+			content: 'content_3',
+			customButtons: [],
+			enabled: true,
+			description: undefined
+		};
+
+		// If I open the wizard and then add a page
+		mainThreadModelViewDialog.$openWizard(wizardHandle);
+		mainThreadModelViewDialog.$setWizardPageDetails(page3Handle, page3Details);
+		mainThreadModelViewDialog.$addWizardPage(wizardHandle, page3Handle, 0);
+
+		// Then the updated page info gets sent to the extension host
+		mockExtHostModelViewDialog.verify(x => x.$updateWizardPageInfo(
+			It.is(handle => handle === wizardHandle),
+			It.is(pageHandles => pageHandles.length === 3 && pageHandles[0] === page3Handle),
+			It.is(currentPage => currentPage === 1)), Times.once());
+	});
+
+	test('removeWizardPage method removes pages at the correct spot and notifies the extension host', () => {
+		mockExtHostModelViewDialog.setup(x => x.$updateWizardPageInfo(It.isAny(), It.isAny(), It.isAny()));
+
+		// If I open the wizard and then remove a page
+		mainThreadModelViewDialog.$openWizard(wizardHandle);
+		mainThreadModelViewDialog.$removeWizardPage(wizardHandle, 0);
+
+		// Then the updated page info gets sent to the extension host
+		mockExtHostModelViewDialog.verify(x => x.$updateWizardPageInfo(
+			It.is(handle => handle === wizardHandle),
+			It.is(pageHandles => pageHandles.length === 1 && pageHandles[0] === page2Handle),
+			It.is(currentPage => currentPage === 0)), Times.once());
+	});
+
+	test('Creating a wizard adds a navigation validation that calls the extension host', () => {
+		mockExtHostModelViewDialog.setup(x => x.$validateNavigation(It.isAny(), It.isAny()));
+
+		// If I call validateNavigation on the wizard that gets created
+		mainThreadModelViewDialog.$openWizard(wizardHandle);
+		openedWizard.validateNavigation(1);
+
+		// Then the call gets forwarded to the extension host
+		mockExtHostModelViewDialog.verify(x => x.$validateNavigation(It.is(handle => handle === wizardHandle), It.is(info => info.newPage === 1)), Times.once());
+	});
+
+	test('Adding a message to a wizard fires events on the created wizard', () => {
+		mainThreadModelViewDialog.$openWizard(wizardHandle);
+		let newMessage: DialogMessage;
+		openedWizard.onMessageChange(message => newMessage = message);
+
+		// If I change the wizard's message
+		wizardDetails.message = {
+			level: MessageLevel.Error,
+			text: 'test message'
+		};
+		mainThreadModelViewDialog.$setWizardDetails(wizardHandle, wizardDetails);
+
+		// Then the message gets changed on the wizard
+		assert.equal(newMessage, wizardDetails.message, 'New message was not included in the fired event');
+		assert.equal(openedWizard.message, wizardDetails.message, 'New message was not set on the wizard');
+	});
+
+	test('Creating a dialog adds a close validation that calls the extension host', () => {
+		mockExtHostModelViewDialog.setup(x => x.$validateDialogClose(It.isAny()));
+
+		// If I call validateClose on the dialog that gets created
+		mainThreadModelViewDialog.$openDialog(dialogHandle);
+		openedDialog.validateClose();
+
+		// Then the call gets forwarded to the extension host
+		mockExtHostModelViewDialog.verify(x => x.$validateDialogClose(It.is(handle => handle === dialogHandle)), Times.once());
+	});
+});

src/sql/workbench/test/electron-browser/api/mainThreadNotebook.test.ts

@@ -0,0 +1,180 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import * as TypeMoq from 'typemoq';
+import * as azdata from 'azdata';
+
+import { URI, UriComponents } from 'vs/base/common/uri';
+import { IExtHostContext } from 'vs/workbench/api/common/extHost.protocol';
+
+import { MainThreadNotebook } from 'sql/workbench/api/browser/mainThreadNotebook';
+import { NotebookService } from 'sql/workbench/services/notebook/common/notebookServiceImpl';
+import { INotebookProvider } from 'sql/workbench/services/notebook/common/notebookService';
+import { INotebookManagerDetails, INotebookSessionDetails, INotebookKernelDetails, INotebookFutureDetails } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { LocalContentManager } from 'sql/workbench/services/notebook/node/localContentManager';
+import { TestLifecycleService } from 'vs/workbench/test/workbenchTestServices';
+import { MockContextKeyService } from 'vs/platform/keybinding/test/common/mockKeybindingService';
+import { ExtHostNotebookShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+
+suite('MainThreadNotebook Tests', () => {
+
+	let mainThreadNotebook: MainThreadNotebook;
+	let mockProxy: TypeMoq.Mock<ExtHostNotebookShape>;
+	let notebookUri: URI;
+	let mockNotebookService: TypeMoq.Mock<NotebookService>;
+	let providerId = 'TestProvider';
+
+	setup(() => {
+		mockProxy = TypeMoq.Mock.ofType(ExtHostNotebookStub);
+		let extContext = <IExtHostContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+		mockNotebookService = TypeMoq.Mock.ofType(NotebookService, undefined, new TestLifecycleService(), undefined, undefined, undefined, undefined, new MockContextKeyService());
+		notebookUri = URI.parse('file:/user/default/my.ipynb');
+		mainThreadNotebook = new MainThreadNotebook(extContext, mockNotebookService.object);
+	});
+
+	suite('On registering a provider', () => {
+		let provider: INotebookProvider;
+		let registeredProviderId: string;
+		setup(() => {
+			mockNotebookService.setup(s => s.registerProvider(TypeMoq.It.isAnyString(), TypeMoq.It.isAny())).returns((id, providerImpl) => {
+				registeredProviderId = id;
+				provider = providerImpl;
+			});
+		});
+
+		test('should call through to notebook service', () => {
+			// When I register a provider
+			mainThreadNotebook.$registerNotebookProvider(providerId, 1);
+			// Then I expect a provider implementation to be passed to the service
+			mockNotebookService.verify(s => s.registerProvider(TypeMoq.It.isAnyString(), TypeMoq.It.isAny()), TypeMoq.Times.once());
+			assert.equal(provider.providerId, providerId);
+		});
+		test('should unregister in service', () => {
+			// Given we have a provider
+			mainThreadNotebook.$registerNotebookProvider(providerId, 1);
+			// When I unregister a provider twice
+			mainThreadNotebook.$unregisterNotebookProvider(1);
+			mainThreadNotebook.$unregisterNotebookProvider(1);
+			// Then I expect it to be unregistered in the service just 1 time
+			mockNotebookService.verify(s => s.unregisterProvider(TypeMoq.It.isValue(providerId)), TypeMoq.Times.once());
+		});
+	});
+
+	suite('getNotebookManager', () => {
+		let managerWithAllFeatures: INotebookManagerDetails;
+		let provider: INotebookProvider;
+
+		setup(() => {
+			managerWithAllFeatures = {
+				handle: 2,
+				hasContentManager: true,
+				hasServerManager: true
+			};
+			mockNotebookService.setup(s => s.registerProvider(TypeMoq.It.isAnyString(), TypeMoq.It.isAny())).returns((id, providerImpl) => {
+				provider = providerImpl;
+			});
+			mainThreadNotebook.$registerNotebookProvider(providerId, 1);
+
+			// Always return empty specs in this test suite
+			mockProxy.setup(p => p.$refreshSpecs(TypeMoq.It.isAnyNumber())).returns(() => Promise.resolve(undefined));
+		});
+
+		test('should return manager with default content manager & undefined server manager if extension host has none', async () => {
+			// Given the extension provider doesn't have acontent or server manager
+			let details: INotebookManagerDetails = {
+				handle: 2,
+				hasContentManager: false,
+				hasServerManager: false
+			};
+			mockProxy.setup(p => p.$getNotebookManager(TypeMoq.It.isAnyNumber(), TypeMoq.It.isValue(notebookUri)))
+				.returns(() => Promise.resolve(details));
+
+			// When I get the notebook manager
+			let manager = await provider.getNotebookManager(notebookUri);
+
+			// Then it should use the built-in content manager
+			assert.ok(manager.contentManager instanceof LocalContentManager);
+			// And it should not define a server manager
+			assert.equal(manager.serverManager, undefined);
+		});
+
+		test('should return manager with a content & server manager if extension host has these', async () => {
+			// Given the extension provider doesn't have acontent or server manager
+			mockProxy.setup(p => p.$getNotebookManager(TypeMoq.It.isAnyNumber(), TypeMoq.It.isValue(notebookUri)))
+				.returns(() => Promise.resolve(managerWithAllFeatures));
+
+			// When I get the notebook manager
+			let manager = await provider.getNotebookManager(notebookUri);
+
+			// Then it shouldn't have wrappers for the content or server manager
+			assert.ok(!(manager.contentManager instanceof LocalContentManager));
+			assert.notEqual(manager.serverManager, undefined);
+		});
+	});
+
+});
+
+class ExtHostNotebookStub implements ExtHostNotebookShape {
+	$getNotebookManager(providerHandle: number, notebookUri: UriComponents): Thenable<INotebookManagerDetails> {
+		throw new Error('Method not implemented.');
+	}
+	$handleNotebookClosed(notebookUri: UriComponents): void {
+		throw new Error('Method not implemented.');
+	}
+	$doStartServer(managerHandle: number): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
+	$doStopServer(managerHandle: number): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
+	$getNotebookContents(managerHandle: number, notebookUri: UriComponents): Thenable<azdata.nb.INotebookContents> {
+		throw new Error('Method not implemented.');
+	}
+	$save(managerHandle: number, notebookUri: UriComponents, notebook: azdata.nb.INotebookContents): Thenable<azdata.nb.INotebookContents> {
+		throw new Error('Method not implemented.');
+	}
+	$refreshSpecs(managerHandle: number): Thenable<azdata.nb.IAllKernels> {
+		throw new Error('Method not implemented.');
+	}
+	$startNewSession(managerHandle: number, options: azdata.nb.ISessionOptions): Thenable<INotebookSessionDetails> {
+		throw new Error('Method not implemented.');
+	}
+	$shutdownSession(managerHandle: number, sessionId: string): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
+	$changeKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable<INotebookKernelDetails> {
+		throw new Error('Method not implemented.');
+	}
+	$configureKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
+	$configureConnection(sessionId: number, conneection: azdata.IConnectionProfile): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
+	$getKernelReadyStatus(kernelId: number): Thenable<azdata.nb.IInfoReply> {
+		throw new Error('Method not implemented.');
+	}
+	$getKernelSpec(kernelId: number): Thenable<azdata.nb.IKernelSpec> {
+		throw new Error('Method not implemented.');
+	}
+	$requestComplete(kernelId: number, content: azdata.nb.ICompleteRequest): Thenable<azdata.nb.ICompleteReplyMsg> {
+		throw new Error('Method not implemented.');
+	}
+	$requestExecute(kernelId: number, content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): Thenable<INotebookFutureDetails> {
+		throw new Error('Method not implemented.');
+	}
+	$interruptKernel(kernelId: number): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
+	$sendInputReply(futureId: number, content: azdata.nb.IInputReply): void {
+		throw new Error('Method not implemented.');
+	}
+	$disposeFuture(futureId: number): void {
+		throw new Error('Method not implemented.');
+	}
+}