[Notebook] Open Kernel Alias (Kusto) notebook with Kernel Alias (Kusto) as selected kernel (#14504)

* Open kernel alias notebook with kernel alias as selected kernel

* Add language info test

* fix test

Co-authored-by: Vasu Bhog <bhogvu@mail.uc.edu>

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

@@ -337,6 +337,7 @@ export class KernelsDropdown extends SelectBox {
 					index = 0;
 				}
 				this.setOptions(kernels, index);
+				this.model.selectedKernelDisplayName = kernels[index];
 			}
 		} else if (this.model.clientSession?.isInErrorState) {
 			kernels.unshift(noKernelName);

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

@@ -69,7 +69,7 @@ let expectedNotebookContentOneCell: nb.INotebookContents = {
 	cells: [{
 		cell_type: CellTypes.Code,
 		source: ['insert into t1 values (c1, c2)'],
-		metadata: { language: 'python' },
+		metadata: { language: 'sql' },
 		execution_count: 1
 	}],
 	metadata: {
@@ -83,6 +83,27 @@ let expectedNotebookContentOneCell: nb.INotebookContents = {
 	nbformat_minor: 5
 };
 
+let expectedKernelAliasNotebookContentOneCell: nb.INotebookContents = {
+	cells: [{
+		cell_type: CellTypes.Code,
+		source: ['StormEvents | summarize Count = count() by State | sort by Count | limit 10'],
+		execution_count: 1
+	}],
+	metadata: {
+		kernelspec: {
+			name: 'mssql',
+			language: 'sql',
+			display_name: 'SQL'
+		},
+		language_info: {
+			name: 'fake',
+			version: ''
+		}
+	},
+	nbformat: 4,
+	nbformat_minor: 5
+};
+
 let expectedParameterizedNotebookContent: nb.INotebookContents = {
 	cells: [{
 		cell_type: CellTypes.Code,
@@ -601,7 +622,7 @@ suite('notebook model', function (): void {
 	});
 
 	test('Should not be in error state if client session initialization succeeds', async function (): Promise<void> {
-		let model = await loadModelAndStartClientSession();
+		let model = await loadModelAndStartClientSession(expectedNotebookContent);
 
 		assert.equal(model.inErrorState, false);
 		assert.equal(model.notebookManagers.length, 1);
@@ -628,7 +649,7 @@ suite('notebook model', function (): void {
 	});
 
 	test('Should close active session when closed', async function () {
-		let model = await loadModelAndStartClientSession();
+		let model = await loadModelAndStartClientSession(expectedNotebookContent);
 		// After client session is started, ensure session is ready
 		assert(model.isSessionReady);
 
@@ -641,7 +662,7 @@ suite('notebook model', function (): void {
 	});
 
 	test('Should disconnect when connection profile created by notebook', async function () {
-		let model = await loadModelAndStartClientSession();
+		let model = await loadModelAndStartClientSession(expectedNotebookContent);
 		// 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);
@@ -657,7 +678,7 @@ suite('notebook model', function (): void {
 	});
 
 	test('Should not disconnect when connection profile not created by notebook', async function () {
-		let model = await loadModelAndStartClientSession();
+		let model = await loadModelAndStartClientSession(expectedNotebookContent);
 		// Ensure notebook prefix isn't present in connection URI
 		queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.default}some/path`);
 		await changeContextWithConnectionProfile(model);
@@ -692,7 +713,7 @@ suite('notebook model', function (): void {
 	});
 
 	test('Should connect to Fake (kernel alias) connection and set kernelAlias', async function () {
-		let model = await loadModelAndStartClientSession();
+		let model = await loadModelAndStartClientSession(expectedNotebookContent);
 
 		// Ensure notebook prefix is present in the connection URI
 		queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`);
@@ -717,7 +738,7 @@ suite('notebook model', function (): void {
 	});
 
 	test('Should change kernel when connecting to a Fake (kernel alias) connection', async function () {
-		let model = await loadModelAndStartClientSession();
+		let model = await loadModelAndStartClientSession(expectedNotebookContent);
 		// Ensure notebook prefix is present in the connection URI
 		queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`);
 		await changeContextWithFakeConnectionProfile(model);
@@ -742,7 +763,7 @@ suite('notebook model', function (): void {
 	});
 
 	test('Should change kernel from Fake (kernel alias) to SQL kernel when connecting to SQL connection', async function () {
-		let model = await loadModelAndStartClientSession();
+		let model = await loadModelAndStartClientSession(expectedNotebookContent);
 
 		// Ensure notebook prefix is present in the connection URI
 		queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`);
@@ -865,9 +886,28 @@ suite('notebook model', function (): void {
 		assert.equal(output.metadata['multi_connection_mode'], true, 'multi_connection_mode not saved correctly to notebook metadata');
 	});
 
-	async function loadModelAndStartClientSession(): Promise<NotebookModel> {
+	test('Should keep kernel alias as language info kernel alias name even if kernel spec is seralized as SQL', async function () {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
-		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedKernelAliasNotebookContentOneCell));
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		queryConnectionService.setup(c => c.getActiveConnections(TypeMoq.It.isAny())).returns(() => null);
+
+		// Given I have a session that fails to start
+		sessionReady.resolve();
+
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		await model.requestModelLoad();
+
+		// Check to see if language info is set to kernel alias
+		assert.equal(model.languageInfo.name, 'fake', 'Notebook language info is not set properly');
+	});
+
+	async function loadModelAndStartClientSession(notebookContent: nb.INotebookContents): Promise<NotebookModel> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(notebookContent));
 		defaultModelOptions.contentManager = mockContentManager.object;
 
 		queryConnectionService.setup(c => c.getActiveConnections(TypeMoq.It.isAny())).returns(() => null);
@@ -880,6 +920,7 @@ suite('notebook model', function (): void {
 		});
 		let model = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, capabilitiesService);
 		model.onClientSessionReady((session) => actualSession = session);
+
 		await model.requestModelLoad();
 
 		await model.startSession(notebookManagers[0]);

src/sql/workbench/services/notebook/browser/models/notebookModel.ts

@@ -288,6 +288,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._selectedKernelDisplayName;
 	}
 
+	public set selectedKernelDisplayName(kernel: string) {
+		this._selectedKernelDisplayName = kernel;
+	}
+
 	public set trustedMode(isTrusted: boolean) {
 		this._trustedMode = isTrusted;
 
@@ -951,6 +955,11 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	private async updateKernelInfoOnKernelChange(kernel: nb.IKernel, kernelAlias?: string) {
 		await this.updateKernelInfo(kernel);
+		this.kernelAliases.forEach(kernel => {
+			if (this._defaultLanguageInfo?.name === kernel.toLowerCase()) {
+				kernelAlias = kernel;
+			}
+		});
 		if (kernel.info) {
 			this.updateLanguageInfo(kernel.info.language_info);
 		}