diff --git a/test/automation/src/sql/managePackagesDialog.ts b/test/automation/src/sql/managePackagesDialog.ts
index 0986876677..283426a319 100644
--- a/test/automation/src/sql/managePackagesDialog.ts
+++ b/test/automation/src/sql/managePackagesDialog.ts
@@ -19,7 +19,7 @@ export class ManagePackagesDialog extends Dialog {
 		await this.waitForNewDialog();
 	}
 
-	async addNewPackage(packageName: string, packageVersion: string = ''): Promise<void> {
+	async addNewPackage(packageName: string): Promise<string> {
 		const addNewTab = `${ManagePackagesDialog.dialogPage} div[class="tab-header"][aria-controls="dialogPane.Manage Packages.1"]`;
 		await this.code.waitAndClick(addNewTab);
 
@@ -37,12 +37,12 @@ export class ManagePackagesDialog extends Dialog {
 		const searchButton = `${ManagePackagesDialog.dialogPage} a[class="monaco-button monaco-text-button"][aria-label="Search"][aria-disabled="false"]`;
 		await this.code.waitAndClick(searchButton);
 
-		if (packageVersion) {
-			const versionSelectBox = `${ManagePackagesDialog.dialogPage} select[class="monaco-select-box monaco-select-box-dropdown-padding"][aria-label="Package Version"]`;
-			const versionOption = `${versionSelectBox} option[value="${packageVersion}]`;
-			this.code.waitForElement(versionOption);
-			this.code.waitForSetValue(versionSelectBox, packageVersion);
-		}
+		const packageNameSelector = `${ManagePackagesDialog.dialogPage} div[id="textContainer"] span`;
+		await this.code.waitForTextContent(packageNameSelector, packageName);
+
+		// Get the latest package version
+		const versionSelectBox = `${ManagePackagesDialog.dialogPage} select[aria-label="Package Version"] option`;
+		let packageVersion = await this.code.waitForTextContent(versionSelectBox);
 
 		const installButton = `${ManagePackagesDialog.dialogPage} a[class="monaco-button monaco-text-button"][aria-label="Install"][aria-disabled="false"]`;
 		await this.code.waitAndClick(installButton);
@@ -57,5 +57,7 @@ export class ManagePackagesDialog extends Dialog {
 		const closeButton = '.modal .modal-footer a[class="monaco-button monaco-text-button"][aria-label="Close"][aria-disabled="false"]';
 		await this.code.waitAndClick(closeButton);
 		await this.waitForDialogGone();
+
+		return packageVersion;
 	}
 }
diff --git a/test/smoke/src/sql/areas/notebook/notebook.test.ts b/test/smoke/src/sql/areas/notebook/notebook.test.ts
index 856abea9cc..52622760ae 100644
--- a/test/smoke/src/sql/areas/notebook/notebook.test.ts
+++ b/test/smoke/src/sql/areas/notebook/notebook.test.ts
@@ -86,8 +86,7 @@ export function setup(opts: minimist.ParsedArgs) {
 			await app.workbench.sqlNotebook.waitForActiveCellResults();
 		});
 
-		// Skip this test for now since it is not stable - the way that the wizard is initialized needs to be refactored
-		it.skip('can add a new package from the Manage Packages wizard', async function () {
+		it('can add a new package from the Manage Packages wizard', async function () {
 			const app = this.app as Application;
 			await app.workbench.sqlNotebook.newUntitledNotebook();
 			await app.workbench.sqlNotebook.notebookToolbar.waitForKernel('SQL');
@@ -101,9 +100,9 @@ export function setup(opts: minimist.ParsedArgs) {
 
 			await app.workbench.sqlNotebook.notebookToolbar.managePackages();
 			await app.workbench.managePackagesDialog.waitForManagePackagesDialog();
-			await app.workbench.managePackagesDialog.addNewPackage('pyarrow', '7.0.0');
+			let packageVersion = await app.workbench.managePackagesDialog.addNewPackage('pyarrow');
 			await app.workbench.taskPanel.showTaskPanel();
-			await app.workbench.taskPanel.waitForTaskComplete('Installing pyarrow 7.0.0 succeeded');
+			await app.workbench.taskPanel.waitForTaskComplete(`Installing pyarrow ${packageVersion} succeeded`);
 
 			// There should be no error output when running the cell after pyarrow has been installed
 			await app.workbench.sqlNotebook.runActiveCell();