Fix some of the flaky tests (#11979)

* Fixes to some of the flaky test

* Misc

* fixes

src/kernels/execution/cellExecutionCreator.ts

@@ -8,7 +8,8 @@ import {
     NotebookCell,
     NotebookCellExecution,
     NotebookCellOutput,
-    NotebookCellOutputItem
+    NotebookCellOutputItem,
+    TextDocument
 } from 'vscode';
 import { IKernelController } from '../types';
 
@@ -80,10 +81,10 @@ export class NotebookCellExecutionWrapper implements NotebookCellExecution {
  * Class for mapping cells to an instance of a NotebookCellExecution object
  */
 export class CellExecutionCreator {
-    private static _map = new Map<string, NotebookCellExecutionWrapper>();
+    private static _map = new WeakMap<TextDocument, NotebookCellExecutionWrapper>();
     static getOrCreate(cell: NotebookCell, controller: IKernelController) {
         let cellExecution: NotebookCellExecutionWrapper | undefined;
-        const key = cell.document.uri.toString();
+        const key = cell.document;
         cellExecution = this.get(cell);
         if (!cellExecution) {
             cellExecution = CellExecutionCreator.create(key, cell, controller);
@@ -106,11 +107,11 @@ export class CellExecutionCreator {
         return cellExecution;
     }
     static get(cell: NotebookCell) {
-        const key = cell.document.uri.toString();
+        const key = cell.document;
         return CellExecutionCreator._map.get(key);
     }
 
-    private static create(key: string, cell: NotebookCell, controller: IKernelController) {
+    private static create(key: TextDocument, cell: NotebookCell, controller: IKernelController) {
         const result = new NotebookCellExecutionWrapper(
             controller.createNotebookCellExecution(cell),
             controller.id,

src/kernels/kernelCrashMonitor.unit.test.ts

@@ -0,0 +1,166 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import * as fakeTimers from '@sinonjs/fake-timers';
+import { KernelMessage } from '@jupyterlab/services';
+import { anything, instance, mock, verify, when } from 'ts-mockito';
+import { Disposable, EventEmitter, NotebookCell } from 'vscode';
+import { IApplicationShell } from '../platform/common/application/types';
+import { disposeAllDisposables } from '../platform/common/helpers';
+import { IDisposable } from '../platform/common/types';
+import { createKernelController, TestNotebookDocument } from '../test/datascience/notebook/executionHelper';
+import { KernelCrashMonitor } from './kernelCrashMonitor';
+import {
+    IKernel,
+    IKernelConnectionSession,
+    IKernelController,
+    IKernelProvider,
+    INotebookKernelExecution,
+    LocalKernelSpecConnectionMetadata,
+    RemoteKernelSpecConnectionMetadata
+} from './types';
+import { assert } from 'chai';
+import { DataScience } from '../platform/common/utils/localize';
+import { createOutputWithErrorMessageForDisplay } from '../platform/errors/errorUtils';
+import { getDisplayNameOrNameOfKernelConnection } from './helpers';
+
+suite('Kernel Crash Monitor', () => {
+    let kernelProvider: IKernelProvider;
+    const disposables: IDisposable[] = [];
+    let kernel: IKernel;
+    let appShell: IApplicationShell;
+    let kernelCrashMonitor: KernelCrashMonitor;
+    let onKernelStatusChanged: EventEmitter<{
+        status: KernelMessage.Status;
+        kernel: IKernel;
+    }>;
+    let onDidStartKernel: EventEmitter<IKernel>;
+    let kernelExecution: INotebookKernelExecution;
+    let onPreExecute: EventEmitter<NotebookCell>;
+    let cell: NotebookCell;
+    let kernelSession: IKernelConnectionSession;
+    let notebook: TestNotebookDocument;
+    let controller: IKernelController;
+    let clock: fakeTimers.InstalledClock;
+    let remoteKernelSpec = RemoteKernelSpecConnectionMetadata.create({
+        id: 'remote',
+        baseUrl: '1',
+        kernelSpec: {
+            argv: [],
+            display_name: 'remote',
+            executable: '',
+            name: 'remote'
+        },
+        serverId: '1'
+    });
+    let localKernelSpec = LocalKernelSpecConnectionMetadata.create({
+        id: 'local',
+        kernelSpec: {
+            argv: [],
+            display_name: 'remote',
+            executable: '',
+            name: 'remote'
+        }
+    });
+    setup(async () => {
+        kernelProvider = mock<IKernelProvider>();
+        kernel = mock<IKernel>();
+        appShell = mock<IApplicationShell>();
+        kernelExecution = mock<INotebookKernelExecution>();
+        kernelSession = mock<IKernelConnectionSession>();
+        onKernelStatusChanged = new EventEmitter<{
+            status: KernelMessage.Status;
+            kernel: IKernel;
+        }>();
+        onDidStartKernel = new EventEmitter<IKernel>();
+        onPreExecute = new EventEmitter<NotebookCell>();
+        notebook = new TestNotebookDocument();
+        cell = await notebook.appendCodeCell('1234');
+        controller = createKernelController('1');
+        disposables.push(onDidStartKernel);
+        disposables.push(onKernelStatusChanged);
+        disposables.push(onPreExecute);
+        when(kernel.dispose()).thenResolve();
+        when(kernel.disposed).thenReturn(false);
+        when(kernel.controller).thenReturn(controller);
+        when(kernel.disposing).thenReturn(false);
+        when(kernel.session).thenReturn(instance(kernelSession));
+        when(kernelSession.kind).thenReturn('localRaw');
+        when(appShell.showErrorMessage(anything())).thenResolve();
+
+        when(kernelProvider.onDidStartKernel).thenReturn(onDidStartKernel.event);
+        when(kernelProvider.onKernelStatusChanged).thenReturn(onKernelStatusChanged.event);
+        when(kernelProvider.getOrCreate(anything(), anything())).thenReturn(instance(kernel));
+        when(kernelProvider.getKernelExecution(anything())).thenReturn(instance(kernelExecution));
+        when(kernelExecution.onPreExecute).thenReturn(onPreExecute.event);
+
+        kernelCrashMonitor = new KernelCrashMonitor(disposables, instance(appShell), instance(kernelProvider));
+        clock = fakeTimers.install();
+        disposables.push(new Disposable(() => clock.uninstall()));
+    });
+    teardown(() => disposeAllDisposables(disposables));
+
+    test('Error message displayed and Cell output updated with error message (raw kernel)', async () => {
+        when(kernelSession.kind).thenReturn('localRaw');
+        when(kernel.kernelConnectionMetadata).thenReturn(localKernelSpec);
+        // Ensure we have a kernel and have started a cell.
+        kernelCrashMonitor.activate();
+        onDidStartKernel.fire(instance(kernel));
+        onPreExecute.fire(cell);
+        const execution = controller.createNotebookCellExecution(cell);
+        execution.start();
+
+        const expectedErrorMessage = Buffer.from(
+            createOutputWithErrorMessageForDisplay(DataScience.kernelCrashedDueToCodeInCurrentOrPreviousCell())
+                ?.items[0]!.data!
+        ).toString();
+
+        when(kernel.status).thenReturn('dead');
+        onKernelStatusChanged.fire({ status: 'dead', kernel: instance(kernel) });
+        await clock.runAllAsync();
+
+        verify(
+            appShell.showErrorMessage(
+                DataScience.kernelDiedWithoutError().format(getDisplayNameOrNameOfKernelConnection(localKernelSpec))
+            )
+        ).once();
+
+        assert.strictEqual(cell.outputs.length, 1);
+        assert.strictEqual(cell.outputs[0].items.length, 1);
+        const outputItem = cell.outputs[0].items[0];
+        assert.include(Buffer.from(outputItem.data).toString(), expectedErrorMessage);
+    });
+    test('Error message displayed and Cell output updated with error message (jupyter kernel)', async () => {
+        when(kernelSession.kind).thenReturn('localJupyter');
+        when(kernel.kernelConnectionMetadata).thenReturn(remoteKernelSpec);
+        // Ensure we have a kernel and have started a cell.
+        kernelCrashMonitor.activate();
+        onDidStartKernel.fire(instance(kernel));
+        onPreExecute.fire(cell);
+        const execution = controller.createNotebookCellExecution(cell);
+        execution.start();
+
+        const expectedErrorMessage = Buffer.from(
+            createOutputWithErrorMessageForDisplay(DataScience.kernelCrashedDueToCodeInCurrentOrPreviousCell())
+                ?.items[0]!.data!
+        ).toString();
+
+        when(kernel.status).thenReturn('autorestarting');
+        when(kernelSession.status).thenReturn('autorestarting');
+        onKernelStatusChanged.fire({ status: 'dead', kernel: instance(kernel) });
+        await clock.runAllAsync();
+
+        verify(
+            appShell.showErrorMessage(
+                DataScience.kernelDiedWithoutErrorAndAutoRestarting().format(
+                    getDisplayNameOrNameOfKernelConnection(remoteKernelSpec)
+                )
+            )
+        ).once();
+
+        assert.strictEqual(cell.outputs.length, 1);
+        assert.strictEqual(cell.outputs[0].items.length, 1);
+        const outputItem = cell.outputs[0].items[0];
+        assert.include(Buffer.from(outputItem.data).toString(), expectedErrorMessage);
+    });
+});