diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
index 01607b91ed5..cfee10b8f04 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
@@ -454,8 +454,9 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 		DOM.toggleClass(this.getDomNode(), 'notebook-editor-editable', !!this.viewModel!.metadata?.editable);
 	}
 
-	private createWebview(id: string, document: URI) {
+	private async createWebview(id: string, document: URI): Promise<void> {
 		this.webview = this.instantiationService.createInstance(BackLayerWebView, this, id, document);
+		await this.webview.waitForInitialization();
 		this.webview.webview.onDidBlur(() => this.updateEditorFocus());
 		this.webview.webview.onDidFocus(() => this.updateEditorFocus());
 		this.localStore.add(this.webview.onMessage(message => {
@@ -467,8 +468,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 	}
 
 	private async attachModel(textModel: NotebookTextModel, viewState: INotebookEditorViewState | undefined) {
-		this.createWebview(this.getId(), textModel.uri);
-		await this.webview!.waitForInitialization();
+		await this.createWebview(this.getId(), textModel.uri);
 
 		this.eventDispatcher = new NotebookEventDispatcher();
 		this.viewModel = this.instantiationService.createInstance(NotebookViewModel, textModel.viewType, textModel, this.eventDispatcher, this.getLayoutInfo());
diff --git a/test/automation/src/notebook.ts b/test/automation/src/notebook.ts
index c0b2fc5ed7c..6f1b1a10e99 100644
--- a/test/automation/src/notebook.ts
+++ b/test/automation/src/notebook.ts
@@ -79,7 +79,7 @@ export class Notebook {
 
 	async focusInCellOutput(): Promise<void> {
 		await this.quickAccess.runCommand('notebook.cell.focusInOutput');
-		await this.code.waitForActiveElement('webview');
+		await this.code.waitForActiveElement('webview, .webview');
 	}
 
 	async focusOutCellOutput(): Promise<void> {
diff --git a/test/automation/src/playwrightDriver.ts b/test/automation/src/playwrightDriver.ts
index ad0ab8f5b4b..9d5bb04af6d 100644
--- a/test/automation/src/playwrightDriver.ts
+++ b/test/automation/src/playwrightDriver.ts
@@ -25,7 +25,8 @@ const vscodeToPlaywrightKey: { [key: string]: string } = {
 	up: 'ArrowUp',
 	down: 'ArrowDown',
 	left: 'ArrowLeft',
-	home: 'Home'
+	home: 'Home',
+	esc: 'Escape'
 };
 
 function buildDriver(browser: playwright.Browser, page: playwright.Page): IDriver {