diff --git a/test/automation/src/sql/notebook.ts b/test/automation/src/sql/notebook.ts
index e8923e35e4..c4fe7156ff 100644
--- a/test/automation/src/sql/notebook.ts
+++ b/test/automation/src/sql/notebook.ts
@@ -133,6 +133,22 @@ export class Notebook {
 		});
 		await this.waitForResultsGone(cellIds);
 	}
+
+	async waitForTrustedElements(): Promise<void> {
+		const cellSelector = '.notebookEditor .notebook-cell';
+		await this.code.waitForElement(`${cellSelector} iframe`);
+		await this.code.waitForElement(`${cellSelector} dialog`);
+		await this.code.waitForElement(`${cellSelector} embed`);
+		await this.code.waitForElement(`${cellSelector} svg`);
+	}
+
+	async waitForTrustedElementsGone(): Promise<void> {
+		const cellSelector = '.notebookEditor .notebook-cell';
+		await this.code.waitForElementGone(`${cellSelector} iframe`);
+		await this.code.waitForElementGone(`${cellSelector} dialog`);
+		await this.code.waitForElementGone(`${cellSelector} embed`);
+		await this.code.waitForElementGone(`${cellSelector} svg`);
+	}
 }
 
 export class NotebookToolbar {
diff --git a/test/smoke/src/sql/areas/notebook/notebook.test.ts b/test/smoke/src/sql/areas/notebook/notebook.test.ts
index 666fc96f4c..cd174d1879 100644
--- a/test/smoke/src/sql/areas/notebook/notebook.test.ts
+++ b/test/smoke/src/sql/areas/notebook/notebook.test.ts
@@ -46,13 +46,19 @@ export function setup() {
 			await app.workbench.sqlNotebook.openFile('untrusted.ipynb');
 			await app.workbench.sqlNotebook.waitForKernel('SQL');
 			await app.workbench.sqlNotebook.waitForNotTrustedIcon();
+			await app.workbench.sqlNotebook.waitForTrustedElementsGone();
+
 			await app.workbench.sqlNotebook.trustNotebook();
 			await app.workbench.sqlNotebook.waitForTrustedIcon();
+			await app.workbench.sqlNotebook.waitForTrustedElements();
+
 			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.sqlNotebook.waitForTrustedElements();
+
 			await app.workbench.quickaccess.runCommand('workbench.action.closeActiveEditor');
 		});
 	});