From 0e84dd240b35a5373c852022a649115bf109d6d2 Mon Sep 17 00:00:00 2001
From: Chris LaFreniere <40371649+chlafreniere@users.noreply.github.com>
Date: Mon, 1 Feb 2021 14:21:24 -0800
Subject: [PATCH] Notebook Cell Conversion Language Fixes (#14120)

* Cell language override fixes

* PR feedback
---
 .../electron-browser/notebookModel.test.ts    | 34 +++++++++++++++++--
 .../notebook/browser/models/notebookModel.ts  |  2 ++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts
index 3415e05659..40cebf0c13 100644
--- a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts
+++ b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts
@@ -44,12 +44,11 @@ let expectedNotebookContent: nb.INotebookContents = {
 	cells: [{
 		cell_type: CellTypes.Code,
 		source: ['insert into t1 values (c1, c2)'],
-		metadata: { language: 'python' },
+		metadata: { language: 'sql' },
 		execution_count: 1
 	}, {
 		cell_type: CellTypes.Markdown,
 		source: ['I am *markdown*'],
-		metadata: { language: 'python' },
 		execution_count: 1
 	}],
 	metadata: {
@@ -57,6 +56,9 @@ let expectedNotebookContent: nb.INotebookContents = {
 			name: 'mssql',
 			language: 'sql',
 			display_name: 'SQL'
+		},
+		language_info: {
+			name: 'sql'
 		}
 	},
 	nbformat: 4,
@@ -544,6 +546,34 @@ suite('notebook model', function (): void {
 		assert.equal(notebookContentChange.changeType, NotebookChangeType.CellMetadataUpdated, 'notebookContentChange changeType should indicate ');
 	});
 
+	test('Should set cell language correctly after cell type conversion', async function (): Promise<void> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		let newCell: ICellModel;
+		model.onCellTypeChanged(c => newCell = c);
+
+		let firstCell = model.cells[0];
+		let secondCell = model.cells[1];
+
+		assert.equal(firstCell.cellType, CellTypes.Code, 'Initial cell type for first cell should be code');
+		assert.equal(firstCell.language, 'sql', 'Initial language should be sql for first cell');
+
+		model.convertCellType(firstCell);
+		assert.equal(firstCell.cellType, CellTypes.Markdown, 'Failed to convert cell type after conversion');
+		assert.equal(firstCell.language, 'markdown', 'Language should be markdown for text cells');
+		assert.deepEqual(newCell, firstCell);
+
+		model.convertCellType(secondCell);
+		assert.equal(secondCell.cellType, CellTypes.Code, 'Failed to convert second cell type');
+		assert.equal(secondCell.language, 'sql', 'Language should be sql again for second cell');
+		assert.deepEqual(newCell, secondCell);
+	});
+
 	test('Should load contents but then go to error state if client session startup fails', async function (): Promise<void> {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContentOneCell));
diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts
index d29976f62f..268e3468fd 100644
--- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts
+++ b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts
@@ -581,6 +581,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		if (cell) {
 			let index = this.findCellIndex(cell);
 			if (index > -1) {
+				// Ensure override language is reset
+				cell.setOverrideLanguage('');
 				cell.cellType = cell.cellType === CellTypes.Markdown ? CellTypes.Code : CellTypes.Markdown;
 				this._onCellTypeChanged.fire(cell);
 				this._contentChangedEmitter.fire({