From 7cf7ca5d15d0c27165a3758a1b11036a336d0361 Mon Sep 17 00:00:00 2001
From: Vasu Bhog <vabhog@microsoft.com>
Date: Wed, 30 Jun 2021 21:33:54 -0700
Subject: [PATCH] use helper function for smoke test (#15964)

---
 test/automation/src/sql/notebook.ts | 14 ++------------
 1 file changed, 2 insertions(+), 12 deletions(-)

diff --git a/test/automation/src/sql/notebook.ts b/test/automation/src/sql/notebook.ts
index fa04da284e..729688050f 100644
--- a/test/automation/src/sql/notebook.ts
+++ b/test/automation/src/sql/notebook.ts
@@ -154,12 +154,7 @@ export class Notebook {
 	}
 
 	async waitForAllResults(): Promise<void> {
-		let cellIds: string[] = [];
-		await this.code.waitForElements('div.notebook-cell', false, result => {
-			cellIds = result.map(cell => cell.attributes['id']);
-			return true;
-		});
-		await this.waitForResults(cellIds);
+		await this.waitForResults(await this.getCellIds());
 	}
 
 	async waitForActiveCellResultsGone(): Promise<void> {
@@ -174,12 +169,7 @@ export class Notebook {
 	}
 
 	async waitForAllResultsGone(): Promise<void> {
-		let cellIds: string[] = [];
-		await this.code.waitForElements('div.notebook-cell', false, result => {
-			cellIds = result.map(cell => cell.attributes['id']);
-			return true;
-		});
-		await this.waitForResultsGone(cellIds);
+		await this.waitForResultsGone(await this.getCellIds());
 	}
 
 	async waitForTrustedElements(): Promise<void> {