Fix restart kernel tests

Author: DonJayamanne

src/notebooks/notebookCommandListener.ts

@@ -79,9 +79,9 @@ export class NotebookCommandListener implements IDataScienceCommandListener {
                 Commands.RestartKernel,
                 (context?: { notebookEditor: { notebookUri: Uri } } | Uri) => {
                     if (context && 'notebookEditor' in context) {
-                        this.restartKernel(context?.notebookEditor?.notebookUri).catch(noop);
+                        return this.restartKernel(context?.notebookEditor?.notebookUri).catch(noop);
                     } else {
-                        this.restartKernel(context).catch(noop);
+                        return this.restartKernel(context).catch(noop);
                     }
                 }
             )

src/test/datascience/jupyter/kernels/installationPrompts.vscode.test.ts

@@ -68,6 +68,7 @@ import { isUri } from '../../../../platform/common/utils/misc';
 import { hasErrorOutput, translateCellErrorOutput } from '../../../../kernels/execution/helpers';
 import { BaseKernelError } from '../../../../kernels/errors/types';
 import { IControllerRegistration } from '../../../../notebooks/controllers/types';
+import { traceVerbose } from '../../../../platform/logging';
 
 /* eslint-disable no-invalid-this, , , @typescript-eslint/no-explicit-any */
 suite('Install IPyKernel (install) @kernelInstall', function () {
@@ -562,7 +563,7 @@ suite('Install IPyKernel (install) @kernelInstall', function () {
         await openNotebookAndInstallIpyKernelWhenRunningCell(venvNoKernelPath, venvNoRegPath, 'DoNotInstallIPyKernel');
     });
     // https://github.com/microsoft/vscode-jupyter/issues/12766
-    test.skip('Should be prompted to re-install ipykernel when restarting a kernel from which ipykernel was uninstalled (VSCode Notebook)', async function () {
+    test('Should be prompted to re-install ipykernel when restarting a kernel from which ipykernel was uninstalled (VSCode Notebook)', async function () {
         if (IS_REMOTE_NATIVE_TEST()) {
             return this.skip();
         }
@@ -578,18 +579,27 @@ suite('Install IPyKernel (install) @kernelInstall', function () {
 
         // Confirm message is displayed.
         installerSpy = sinon.spy(installer, 'install');
-        const promptToInstall = await clickInstallFromIPyKernelPrompt();
-        await commandManager.executeCommand(Commands.RestartKernel, {
-            notebookEditor: { notebookUri: nbFile }
-        }),
-            await Promise.all([
-                waitForCondition(
-                    async () => promptToInstall.displayed.then(() => true),
-                    delayForUITest,
-                    'Prompt not displayed'
-                ),
-                waitForIPyKernelToGetInstalled()
-            ]);
+        // Confirm message is displayed.
+        let installedIPyKernel = false;
+
+        // First hookup the prompt for restarting the kernel;
+        const restartPrompt = await clickOKForRestartPrompt();
+        restartPrompt.displayed
+            .then(async () => {
+                restartPrompt.dispose();
+                // Next  hookup the prompt for re-install ipykernel.
+                const selectionPrompt = await clickInstallFromIPyKernelPrompt();
+                installedIPyKernel = await selectionPrompt.displayed;
+            })
+            .catch(noop);
+
+        await Promise.all([
+            commandManager.executeCommand(Commands.RestartKernel, {
+                notebookEditor: { notebookUri: nbFile }
+            }),
+            waitForCondition(async () => installedIPyKernel, delayForUITest, 'Prompt not displayed'),
+            waitForIPyKernelToGetInstalled()
+        ]);
     });
     test('Ensure ipykernel install prompt is displayed and we can select another kernel after uninstalling IPyKernel from a live notebook and then restarting the kernel (VSCode Notebook)', async function () {
         if (IS_REMOTE_NATIVE_TEST()) {
@@ -607,20 +617,29 @@ suite('Install IPyKernel (install) @kernelInstall', function () {
         assert.ok(startController);
 
         // Confirm message is displayed.
-        const promptToInstall = await selectKernelFromIPyKernelPrompt();
-        const { kernelSelected } = hookupKernelSelected(promptToInstall, venvNoRegPath);
-        await commands.executeCommand(Commands.RestartKernel, nbFile);
-
+        let promptToInstallDisplayed = false;
+        let kernelSelected = false;
+
+        // First hookup the prompt for restarting the kernel;
+        const restartPrompt = await clickOKForRestartPrompt();
+        restartPrompt.displayed
+            .then(async () => {
+                restartPrompt.dispose();
+                // Next  hookup the prompt for re-install ipykernel.
+                const selectionPrompt = await selectKernelFromIPyKernelPrompt();
+                traceVerbose('Hooked prompt to change kernel');
+                promptToInstallDisplayed = await selectionPrompt.displayed;
+                traceVerbose('Displayed and handled prompt to change kernel', promptToInstallDisplayed);
+                // Next  hookup the prompt for selecting a kernel
+                kernelSelected = await hookupKernelSelected(restartPrompt, venvNoRegPath).kernelSelected;
+                traceVerbose('Hooked prompt to select a new kernel', kernelSelected);
+            })
+            .catch(noop);
+        traceVerbose('Waiting For all');
         await Promise.all([
-            await commandManager.executeCommand(Commands.RestartKernel, {
-                notebookEditor: { notebookUri: nbFile }
-            }),
-            waitForCondition(
-                async () => promptToInstall.displayed.then(() => true),
-                delayForUITest,
-                'Prompt not displayed'
-            ),
-            waitForCondition(async () => kernelSelected, delayForUITest, 'New kernel not selected'),
+            commandManager.executeCommand(Commands.RestartKernel, { notebookEditor: { notebookUri: nbFile } }),
+            waitForCondition(() => promptToInstallDisplayed, delayForUITest, 'Prompt not displayed'),
+            waitForCondition(() => kernelSelected, delayForUITest, 'New kernel not selected'),
             // Verify the new kernel associated with this notebook is different.
             waitForCondition(
                 async () => {
@@ -756,6 +775,14 @@ suite('Install IPyKernel (install) @kernelInstall', function () {
             disposables
         );
     }
+    async function clickOKForRestartPrompt() {
+        return hijackPrompt(
+            'showInformationMessage',
+            { contains: DataScience.restartKernelMessage },
+            { result: DataScience.restartKernelMessageYes, clickImmediately: true },
+            disposables
+        );
+    }
     async function clickInstallFromIPyKernelPrompt() {
         return hijackPrompt(
             'showInformationMessage',