Add Notebook tests for Kernel Alias connections (#12722)

* More Generic tests for kernel alias connections

* Kernel Alias tests for the Notebook Model

* Updated titles of tests

src/sql/platform/capabilities/test/common/testCapabilitiesService.ts

@@ -16,7 +16,7 @@ import { IDisposable } from 'vs/base/common/lifecycle';
 export class TestCapabilitiesService implements ICapabilitiesService {
 
 	private pgsqlProviderName = 'PGSQL';
-	private kustoProviderName = 'KUSTO';
+	private fakeProviderName = 'FAKE';
 	public _serviceBrand: undefined;
 
 	public capabilities: { [id: string]: ProviderFeatures } = {};
@@ -107,15 +107,15 @@ export class TestCapabilitiesService implements ICapabilitiesService {
 			displayName: 'PostgreSQL',
 			connectionOptions: connectionProvider,
 		};
-		let kustoCapabilities = {
+		let fakeCapabilities = {
-			providerId: this.kustoProviderName,
+			providerId: this.fakeProviderName,
-			displayName: 'Kusto',
+			displayName: 'fakeName',
 			connectionOptions: connectionProvider,
-			notebookKernelAlias: 'Kusto'
+			notebookKernelAlias: 'fakeAlias'
 		};
 		this.capabilities[mssqlProviderName] = { connection: msSQLCapabilities };
 		this.capabilities[this.pgsqlProviderName] = { connection: pgSQLCapabilities };
-		this.capabilities[this.kustoProviderName] = { connection: kustoCapabilities };
+		this.capabilities[this.fakeProviderName] = { connection: fakeCapabilities };
 	}
 
 	registerConnectionProvider(id: string, properties: ConnectionProviderProperties): IDisposable {

src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts

@@ -20,7 +20,6 @@ import { ClientSession } from 'sql/workbench/services/notebook/browser/models/cl
 import { CellTypes, NotebookChangeType } from 'sql/workbench/services/notebook/common/contracts';
 import { Deferred } from 'sql/base/common/promise';
 import { Memento } from 'vs/workbench/common/memento';
-import { Emitter } from 'vs/base/common/event';
 import { TestCapabilitiesService } from 'sql/platform/capabilities/test/common/testCapabilitiesService';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 import { TestStorageService } from 'vs/workbench/test/common/workbenchTestServices';
@@ -79,12 +78,12 @@ let expectedNotebookContentOneCell: nb.INotebookContents = {
 
 let defaultUri = URI.file('/some/path.ipynb');
 
-let mockClientSession: TypeMoq.Mock<IClientSession>;
+let mockClientSession: IClientSession;
 let clientSessionOptions: IClientSessionOptions;
 let sessionReady: Deferred<void>;
 let mockModelFactory: TypeMoq.Mock<ModelFactory>;
 let notificationService: TypeMoq.Mock<INotificationService>;
-let capabilitiesService: TypeMoq.Mock<ICapabilitiesService>;
+let capabilitiesService: ICapabilitiesService;
 let instantiationService: IInstantiationService;
 
 suite('notebook model', function (): void {
@@ -99,7 +98,7 @@ suite('notebook model', function (): void {
 		notebookManagers[0].sessionManager = mockSessionManager.object;
 		sessionReady = new Deferred<void>();
 		notificationService = TypeMoq.Mock.ofType(TestNotificationService, TypeMoq.MockBehavior.Loose);
-		capabilitiesService = TypeMoq.Mock.ofType(TestCapabilitiesService);
+		capabilitiesService = new TestCapabilitiesService();
 		memento = TypeMoq.Mock.ofType(Memento, TypeMoq.MockBehavior.Loose, '');
 		memento.setup(x => x.getMemento(TypeMoq.It.isAny())).returns(() => void 0);
 		queryConnectionService = TypeMoq.Mock.ofType(TestConnectionManagementService, TypeMoq.MockBehavior.Loose, memento.object, undefined, new TestStorageService());
@@ -117,7 +116,7 @@ suite('notebook model', function (): void {
 			cellMagicMapper: undefined,
 			defaultKernel: undefined,
 			layoutChanged: undefined,
-			capabilitiesService: capabilitiesService.object
+			capabilitiesService: capabilitiesService
 		};
 		clientSessionOptions = {
 			notebookManager: defaultModelOptions.notebookManagers[0],
@@ -125,15 +124,12 @@ suite('notebook model', function (): void {
 			notificationService: notificationService.object,
 			kernelSpec: defaultModelOptions.defaultKernel
 		};
-		mockClientSession = TypeMoq.Mock.ofType(ClientSession, undefined, clientSessionOptions);
+		mockClientSession = new ClientSession(clientSessionOptions);
-		mockClientSession.setup(c => c.initialize()).returns(() => {
+		mockClientSession.initialize();
-			return Promise.resolve();
-		});
-		mockClientSession.setup(c => c.ready).returns(() => sessionReady.promise);
 		mockModelFactory = TypeMoq.Mock.ofType(ModelFactory);
 		mockModelFactory.callBase = true;
 		mockModelFactory.setup(f => f.createClientSession(TypeMoq.It.isAny())).returns(() => {
-			return mockClientSession.object;
+			return mockClientSession;
 		});
 	});
 
@@ -373,8 +369,6 @@ suite('notebook model', function (): void {
 		defaultModelOptions.contentManager = mockContentManager.object;
 
 		// Given I have a session that fails to start
-		mockClientSession.setup(c => c.isInErrorState).returns(() => true);
-		mockClientSession.setup(c => c.errorMessage).returns(() => 'Error');
 		sessionReady.resolve();
 		let sessionFired = false;
 
@@ -400,7 +394,7 @@ suite('notebook model', function (): void {
 
 		assert.equal(model.inErrorState, false);
 		assert.equal(model.notebookManagers.length, 1);
-		assert.deepEqual(model.clientSession, mockClientSession.object);
+		assert.deepEqual(model.clientSession, mockClientSession);
 	});
 
 	test('Should notify on trust set', async function () {
@@ -486,16 +480,16 @@ suite('notebook model', function (): void {
 		assert(output.metadata['custom-object']['prop2'] === 'value2', 'Custom metadata for object was not preserved');
 	});
 
-	test('Should get Kusto connection and set kernelAlias', async function () {
+	test('Should connect to Fake (kernel alias) connection and set kernelAlias', async function () {
 		let model = await loadModelAndStartClientSession();
 
 		// Ensure notebook prefix is present in the connection URI
 		queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`);
-		await changeContextWithConnectionProfile(model);
+		await changeContextWithFakeConnectionProfile(model);
 
-		//Check to see if Kusto is added to kernelAliases
+		//Check to see if Alias is added to kernelAliases
 		assert(!isUndefinedOrNull(model.kernelAliases));
-		let expectedAlias = ['Kusto'];
+		let expectedAlias = ['fakeAlias'];
 		let kernelAliases = model.kernelAliases;
 
 		assert.equal(kernelAliases.length, 1);
@@ -511,23 +505,20 @@ suite('notebook model', function (): void {
 		queryConnectionService.verify((c) => c.disconnect(TypeMoq.It.isAny()), TypeMoq.Times.once());
 	});
 
-	test.skip('Should change kernel to Kusto when connecting to Kusto connection', async function () {
+	test('Should change kernel when connecting to a Fake (kernel alias) connection', async function () {
 		let model = await loadModelAndStartClientSession();
-
 		// Ensure notebook prefix is present in the connection URI
 		queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`);
-		await changeContextWithConnectionProfile(model);
+		await changeContextWithFakeConnectionProfile(model);
 
 		// // After client session is started, ensure context isn't null/undefined
 		assert(!isUndefinedOrNull(model.context), 'context should exist after call to change context');
 
+		let notebookKernelAlias = model.context.serverCapabilities.notebookKernelAlias;
 		let doChangeKernelStub = sinon.spy(model, 'doChangeKernel').withArgs(model.kernelAliases[0]);
 
-		//Change to kusto kernel
+		model.changeKernel(notebookKernelAlias);
-		//TODO issue with Test not setting serverCapabilities of context
+		assert.equal(model.selectedKernelDisplayName, notebookKernelAlias);
-		await model.changeKernel(model.kernelAliases[0]);
-		let notebookKernelAlias = capabilitiesService.instance.providers.KUSTO.connection.notebookKernelAlias;
-		assert.equal(model.selectedKernelDisplayName, model.kernelAliases[0]);
 		assert.equal(model.currentKernelAlias, notebookKernelAlias);
 		sinon.assert.called(doChangeKernelStub);
 		sinon.restore(doChangeKernelStub);
@@ -539,17 +530,47 @@ suite('notebook model', function (): void {
 		queryConnectionService.verify((c) => c.disconnect(TypeMoq.It.isAny()), TypeMoq.Times.once());
 	});
 
+	test('Should change kernel from Fake (kernel alias) to SQL kernel when connecting to SQL connection', async function () {
+		let model = await loadModelAndStartClientSession();
+
+		// Ensure notebook prefix is present in the connection URI
+		queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`);
+		// Connect to fake connection enables kernel alias connection
+		await changeContextWithFakeConnectionProfile(model);
+
+		// After client session is started, ensure context isn't null/undefined
+		assert(!isUndefinedOrNull(model.context), 'context should exist after call to change context');
+
+		let notebookKernelAlias = model.context.serverCapabilities.notebookKernelAlias;
+		let doChangeKernelStub = sinon.spy(model, 'doChangeKernel');
+
+		// Change kernel first to alias kernel and then connect to SQL connection
+		model.changeKernel(notebookKernelAlias);
+		assert.equal(model.selectedKernelDisplayName, notebookKernelAlias);
+		assert.equal(model.currentKernelAlias, notebookKernelAlias);
+		sinon.assert.called(doChangeKernelStub);
+		sinon.restore(doChangeKernelStub);
+
+		// Change to SQL connection from Fake connection
+		await changeContextWithConnectionProfile(model);
+		let expectedKernel = 'SQL';
+		model.changeKernel(expectedKernel);
+		assert.equal(model.selectedKernelDisplayName, expectedKernel);
+		assert.equal(model.currentKernelAlias, undefined);
+		sinon.assert.called(doChangeKernelStub);
+		sinon.restore(doChangeKernelStub);
+
+		// After closing the notebook
+		await model.handleClosed();
+
+		// Ensure disconnect is called once
+		queryConnectionService.verify((c) => c.disconnect(TypeMoq.It.isAny()), TypeMoq.Times.once());
+	});
+
 	async function loadModelAndStartClientSession(): Promise<NotebookModel> {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
 		defaultModelOptions.contentManager = mockContentManager.object;
-		let kernelChangedEmitter: Emitter<nb.IKernelChangedArgs> = new Emitter<nb.IKernelChangedArgs>();
-		let statusChangedEmitter: Emitter<nb.ISession> = new Emitter<nb.ISession>();
-
-		mockClientSession.setup(c => c.isInErrorState).returns(() => false);
-		mockClientSession.setup(c => c.isReady).returns(() => true);
-		mockClientSession.setup(c => c.kernelChanged).returns(() => kernelChangedEmitter.event);
-		mockClientSession.setup(c => c.statusChanged).returns(() => statusChangedEmitter.event);
 
 		queryConnectionService.setup(c => c.getActiveConnections(TypeMoq.It.isAny())).returns(() => null);
 
@@ -559,7 +580,6 @@ suite('notebook model', function (): void {
 		let options: INotebookModelOptions = assign({}, defaultModelOptions, <Partial<INotebookModelOptions>>{
 			factory: mockModelFactory.object
 		});
-		let capabilitiesService = new TestCapabilitiesService;
 		let model = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService(), capabilitiesService);
 		model.onClientSessionReady((session) => actualSession = session);
 		await model.requestModelLoad();
@@ -569,7 +589,7 @@ suite('notebook model', function (): void {
 		// Then I expect load to succeed
 		assert(!isUndefinedOrNull(model.clientSession), 'clientSession should exist after session is started');
 
-		assert.deepEqual(actualSession, mockClientSession.object, 'session returned is not the expected object');
+		assert.deepEqual(actualSession, mockClientSession, 'session returned is not the expected object');
 
 		// but on server load completion I expect error state to be set
 		// Note: do not expect serverLoad event to throw even if failed
@@ -578,7 +598,7 @@ suite('notebook model', function (): void {
 	}
 
 	async function changeContextWithConnectionProfile(model: NotebookModel) {
-		let connection = new ConnectionProfile(capabilitiesService.object, {
+		let connection = new ConnectionProfile(capabilitiesService, {
 			connectionName: 'newName',
 			savePassword: false,
 			groupFullName: 'testGroup',
@@ -589,12 +609,32 @@ suite('notebook model', function (): void {
 			userName: 'testUsername',
 			groupId: undefined,
 			providerName: mssqlProviderName,
-			options: { serverCapabilities: capabilitiesService.instance.providers.KUSTO.connection },
+			options: {},
 			saveProfile: true,
 			id: 'testID'
 		});
 
 		await model.changeContext(connection.connectionName, connection);
 	}
+
+	async function changeContextWithFakeConnectionProfile(model: NotebookModel) {
+		let fakeConnection = new ConnectionProfile(capabilitiesService, {
+			connectionName: 'newName',
+			savePassword: false,
+			groupFullName: 'testGroup',
+			serverName: 'testServerName',
+			databaseName: 'testDatabaseName',
+			authenticationType: 'integrated',
+			password: 'test',
+			userName: 'testUsername',
+			groupId: undefined,
+			providerName: 'FAKE',
+			options: {},
+			saveProfile: true,
+			id: 'testID'
+		});
+
+		await model.changeContext(fakeConnection.connectionName, fakeConnection);
+	}
 });
 

src/sql/workbench/services/connection/test/browser/connectionManagementService.test.ts

@@ -1412,7 +1412,7 @@ suite('SQL ConnectionManagementService tests', () => {
 	});
 
 	test('providerNameToDisplayNameMap should return all providers', () => {
-		let expectedNames = ['MSSQL', 'PGSQL', 'KUSTO'];
+		let expectedNames = ['MSSQL', 'PGSQL', 'FAKE'];
 		let providerNames = Object.keys(connectionManagementService.providerNameToDisplayNameMap);
 		assert.equal(providerNames.length, 3);
 		assert.equal(providerNames[0], expectedNames[0]);