Add Notebook Save integration tests (#6103)

* Add Notebook Save integration tests
Wrote test to verify save behavior
Fixed issues found during testing,
specifically around how we notify dirty state change to extensions

* Improved error messages

extensions/integration-tests/src/notebook.test.ts

@@ -10,9 +10,9 @@ import * as assert from 'assert';
 import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { context } from './testContext';
-import { sqlNotebookContent, writeNotebookToFile, sqlKernelMetadata, getFileName, pySparkNotebookContent, pySpark3KernelMetadata, pythonKernelMetadata, sqlNotebookMultipleCellsContent, notebookContentForCellLanguageTest, sqlKernelSpec, pythonKernelSpec, pySpark3KernelSpec } from './notebook.util';
+import { sqlNotebookContent, writeNotebookToFile, sqlKernelMetadata, getFileName, pySparkNotebookContent, pySpark3KernelMetadata, pythonKernelMetadata, sqlNotebookMultipleCellsContent, notebookContentForCellLanguageTest, sqlKernelSpec, pythonKernelSpec, pySpark3KernelSpec, CellTypes } from './notebook.util';
 import { getBdcServer, getConfigValue, EnvironmentVariable_PYTHON_PATH } from './testConfig';
-import { connectToServer } from './utils';
+import { connectToServer, sleep } from './utils';
 import * as fs from 'fs';
 import { stressify } from 'adstest';
 
@@ -44,6 +44,10 @@ if (context.RunTest) {
 			await (new NotebookTester()).sqlLanguageTest(this.test.title);
 		});
 
+		test('should not be dirty after saving notebook test', async function () {
+			await (new NotebookTester().shouldNotBeDirtyAfterSavingNotebookTest(this.test.title));
+		});
+
 		if (process.env['RUN_PYTHON3_TEST'] === '1') {
 			test('Python3 notebook test', async function () {
 				await (new NotebookTester()).python3NbTest(this.test.title);
@@ -180,6 +184,41 @@ class NotebookTester {
 		assert(kernelChanged && notebook.document.kernelSpec.name === 'SQL', `Expected third kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`);
 	}
 
+	async shouldNotBeDirtyAfterSavingNotebookTest(title: string): Promise<void> {
+		// Given a notebook that's been edited (in this case, open notebook runs the 1st cell and adds an output)
+		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title);
+		assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`);
+		assert(notebook.document.kernelSpec.name === 'SQL', `Expected first kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`);
+		assert(notebook.document.isDirty === true, 'Notebook should be dirty after edit');
+
+		// When I save it, it should no longer be dirty
+		let saved = await notebook.document.save();
+		assert(saved === true, 'Expect initial save to succeed');
+		// Note: need to sleep after save as the change events happen after save
+		// We need to give back the thread or the event won't have been drained.
+		// This is consistent with VSCode APIs, so keeping as-is
+		await sleep(10);
+		assert(notebook.document.isDirty === false, 'Notebook should not be dirty after initial save');
+
+		// And when I edit again, should become dirty
+		let edited = await notebook.edit(builder => {
+			builder.insertCell({
+				cell_type: CellTypes.Code,
+				source: ''
+			});
+		});
+		assert(edited === true, 'Expect edit to succeed');
+		await sleep(10);
+		assert(notebook.document.isDirty === true, 'Notebook should be dirty after edit');
+
+		// Finally on 2nd save it should no longer be dirty
+		saved = await notebook.document.save();
+		await sleep(10);
+		assert(saved === true, 'Expect save after edit to succeed');
+		assert(notebook.document.isDirty === false, 'Notebook should not be dirty after 2nd save');
+
+	}
+
 	async pythonChangeKernelDifferentProviderTest(title: string): Promise<void> {
 		let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title);
 		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);

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

@@ -15,6 +15,7 @@ import { IInstantiationService } from 'vs/platform/instantiation/common/instanti
 import { ITextEditorOptions } from 'vs/platform/editor/common/editor';
 import { Schemas } from 'vs/base/common/network';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
+import * as types from 'vs/base/common/types';
 
 import {
 	SqlMainContext, MainThreadNotebookDocumentsAndEditorsShape, SqlExtHostContext, ExtHostNotebookDocumentsAndEditorsShape,
@@ -560,7 +561,7 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		let changeData: INotebookModelChangedData = {
 			// Note: we just send all cells for now, not a diff
 			cells: this.convertCellModelToNotebookCell(editor.cells),
-			isDirty: e.isDirty,
+			isDirty: this.getDirtyState(e, editor),
 			providerId: editor.providerId,
 			providers: editor.providers,
 			uri: editor.uri,
@@ -570,10 +571,16 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		return changeData;
 	}
 
+	private getDirtyState(e: NotebookContentChange, editor: MainThreadNotebookEditor): boolean {
+		if (!types.isUndefinedOrNull(e.isDirty)) {
+			return e.isDirty;
+		}
+		return editor.isDirty;
+	}
+
 	mapChangeKind(changeType: NotebookChangeType): NotebookChangeKind {
 		switch (changeType) {
-			case NotebookChangeType.CellDeleted:
-			case NotebookChangeType.CellsAdded:
+			case NotebookChangeType.CellsModified:
 			case NotebookChangeType.CellOutputUpdated:
 			case NotebookChangeType.CellSourceUpdated:
 			case NotebookChangeType.DirtyStateChanged:

src/sql/workbench/parts/notebook/models/contracts.ts

@@ -37,8 +37,7 @@ export class OutputTypes {
 }
 
 export enum NotebookChangeType {
-	CellsAdded,
-	CellDeleted,
+	CellsModified,
 	CellSourceUpdated,
 	CellOutputUpdated,
 	DirtyStateChanged,

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

@@ -341,7 +341,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		this.updateActiveCell(cell);
 
 		this._contentChangedEmitter.fire({
-			changeType: NotebookChangeType.CellsAdded,
+			changeType: NotebookChangeType.CellsModified,
 			cells: [cell],
 			cellIndex: index
 		});
@@ -375,9 +375,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		if (index > -1) {
 			this._cells.splice(index, 1);
 			this._contentChangedEmitter.fire({
-				changeType: NotebookChangeType.CellDeleted,
+				changeType: NotebookChangeType.CellsModified,
 				cells: [cellModel],
-				cellIndex: index
+				cellIndex: index,
+				isDirty: true
 			});
 		} else {
 			this.notifyError(localize('deleteCellFailed', "Failed to delete cell."));
@@ -401,7 +402,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				this.updateActiveCell(newCells[0]);
 			}
 			this._contentChangedEmitter.fire({
-				changeType: NotebookChangeType.CellsAdded
+				changeType: NotebookChangeType.CellsModified,
+				isDirty: true
 			});
 		}
 	}