From 07df54ee61a6d8baab4dc71e8114ff1286c8dca4 Mon Sep 17 00:00:00 2001
From: Lucy Zhang <luczhan@microsoft.com>
Date: Thu, 12 Nov 2020 06:45:20 -0800
Subject: [PATCH] Notebook: Add smoke test for trusted notebooks (#13330)

* add smoke test for trusted notebooks

* change repo

* update trustNotebook method

* pr comment
---
 test/automation/src/sql/notebook.ts           | 41 +++++++++++++++++++
 .../src/sql/areas/notebook/notebook.test.ts   | 15 +++++++
 2 files changed, 56 insertions(+)

diff --git a/test/automation/src/sql/notebook.ts b/test/automation/src/sql/notebook.ts
index 6872fe378c..dacecbcf86 100644
--- a/test/automation/src/sql/notebook.ts
+++ b/test/automation/src/sql/notebook.ts
@@ -7,6 +7,7 @@ import { Code } from '../code';
 import { QuickAccess } from '../quickaccess';
 import { QuickInput } from '../quickinput';
 import { Editors } from '../editors';
+import { IElement } from '..';
 
 const winOrCtrl = process.platform === 'darwin' ? 'ctrl' : 'win';
 
@@ -64,6 +65,18 @@ export class Notebook {
 		await this.code.waitAndClick(clearResultsButton);
 	}
 
+	async trustNotebook(): Promise<void> {
+		await this.toolbar.trustNotebook();
+	}
+
+	async waitForTrustedIcon(): Promise<void> {
+		await this.toolbar.waitForTrustedIcon();
+	}
+
+	async waitForNotTrustedIcon(): Promise<void> {
+		await this.toolbar.waitForNotTrustedIcon();
+	}
+
 	async waitForTypeInEditor(text: string) {
 		const editor = '.notebook-cell.active .monaco-editor';
 		await this.code.waitAndClick(editor);
@@ -124,6 +137,12 @@ export class Notebook {
 export class NotebookToolbar {
 
 	private static readonly toolbarSelector = '.notebookEditor .editor-toolbar .actions-container';
+	private static readonly toolbarButtonSelector = `${NotebookToolbar.toolbarSelector} a.action-label.codicon.notebook-button.masked-icon`;
+	private static readonly trustedButtonClass = 'action-label codicon notebook-button masked-icon icon-shield';
+	private static readonly trustedButtonSelector = `${NotebookToolbar.toolbarSelector} a[class="${NotebookToolbar.trustedButtonClass}"]`;
+	private static readonly notTrustedButtonClass = 'action-label codicon notebook-button masked-icon icon-shield-x';
+	private static readonly notTrustedButtonSelector = `${NotebookToolbar.toolbarSelector} a[class="${NotebookToolbar.notTrustedButtonClass}"]`;
+
 	constructor(private code: Code) { }
 
 	async changeKernel(kernel: string): Promise<void> {
@@ -136,4 +155,26 @@ export class NotebookToolbar {
 		const kernelDropdownValue = `${NotebookToolbar.toolbarSelector} select[id="kernel-dropdown"][title="${kernel}"]`;
 		await this.code.waitForElement(kernelDropdownValue, undefined, 3000); // wait up to 5 minutes for kernel change
 	}
+
+	async trustNotebook(): Promise<void> {
+		await this.code.waitAndClick(NotebookToolbar.toolbarSelector);
+
+		let buttons: IElement[] = await this.code.waitForElements(NotebookToolbar.toolbarButtonSelector, false);
+		buttons.forEach(async button => {
+			if (button.className.includes('icon-shield-x')) {
+				await this.code.waitAndClick(NotebookToolbar.notTrustedButtonSelector);
+				return;
+			} else if (button.className.includes('icon-shield')) { // notebook is already trusted
+				return;
+			}
+		});
+	}
+
+	async waitForTrustedIcon(): Promise<void> {
+		await this.code.waitForElement(NotebookToolbar.trustedButtonSelector);
+	}
+
+	async waitForNotTrustedIcon(): Promise<void> {
+		await this.code.waitForElement(NotebookToolbar.notTrustedButtonSelector);
+	}
 }
diff --git a/test/smoke/src/sql/areas/notebook/notebook.test.ts b/test/smoke/src/sql/areas/notebook/notebook.test.ts
index 9b53a72db8..666fc96f4c 100644
--- a/test/smoke/src/sql/areas/notebook/notebook.test.ts
+++ b/test/smoke/src/sql/areas/notebook/notebook.test.ts
@@ -40,5 +40,20 @@ export function setup() {
 			await app.workbench.sqlNotebook.waitForKernel('Python 3');
 			await app.workbench.sqlNotebook.waitForAllResults();
 		});
+
+		it('can open untrusted notebook, trust, save, and reopen trusted notebook', async function () {
+			const app = this.app as Application;
+			await app.workbench.sqlNotebook.openFile('untrusted.ipynb');
+			await app.workbench.sqlNotebook.waitForKernel('SQL');
+			await app.workbench.sqlNotebook.waitForNotTrustedIcon();
+			await app.workbench.sqlNotebook.trustNotebook();
+			await app.workbench.sqlNotebook.waitForTrustedIcon();
+			await app.workbench.quickaccess.runCommand('workbench.action.files.save');
+			await app.workbench.quickaccess.runCommand('workbench.action.closeActiveEditor');
+
+			await app.workbench.sqlNotebook.openFile('untrusted.ipynb');
+			await app.workbench.sqlNotebook.waitForTrustedIcon();
+			await app.workbench.quickaccess.runCommand('workbench.action.closeActiveEditor');
+		});
 	});
 }