add test for properly disposing notebook, pythonServer

anthonykim1 · anthonykim1 · commit c0dde023573f · 2025-04-01T11:15:01.000-07:00
diff --git a/src/client/common/vscodeApis/workspaceApis.ts b/src/client/common/vscodeApis/workspaceApis.ts
@@ -60,6 +60,10 @@ export function onDidChangeConfiguration(handler: (e: vscode.ConfigurationChange
     return vscode.workspace.onDidChangeConfiguration(handler);
 }
 
+export function onDidCloseNotebookDocument(handler: (e: vscode.NotebookDocument) => void): vscode.Disposable {
+    return vscode.workspace.onDidCloseNotebookDocument(handler);
+}
+
 export function createFileSystemWatcher(
     globPattern: vscode.GlobPattern,
     ignoreCreateEvents?: boolean,
diff --git a/src/client/repl/nativeRepl.ts b/src/client/repl/nativeRepl.ts
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 // Native Repl class that holds instance of pythonServer and replController
 
 import {
@@ -7,13 +10,12 @@ import {
     QuickPickItem,
     TextEditor,
     Uri,
-    workspace,
     WorkspaceFolder,
 } from 'vscode';
 import { Disposable } from 'vscode-jsonrpc';
 import { PVSC_EXTENSION_ID } from '../common/constants';
 import { showQuickPick } from '../common/vscodeApis/windowApis';
-import { getWorkspaceFolders } from '../common/vscodeApis/workspaceApis';
+import { getWorkspaceFolders, onDidCloseNotebookDocument } from '../common/vscodeApis/workspaceApis';
 import { PythonEnvironment } from '../pythonEnvironments/info';
 import { createPythonServer, PythonServer } from './pythonServer';
 import { executeNotebookCell, openInteractiveREPL, selectNotebookKernel } from './replCommandHandler';
@@ -69,7 +71,7 @@ export class NativeRepl implements Disposable {
      */
     private watchNotebookClosed(): void {
         this.disposables.push(
-            workspace.onDidCloseNotebookDocument(async (nb) => {
+            onDidCloseNotebookDocument(async (nb) => {
                 if (this.notebookDocument && nb.uri.toString() === this.notebookDocument.uri.toString()) {
                     this.notebookDocument = undefined;
                     this.newReplSession = true;
diff --git a/src/test/repl/nativeRepl.test.ts b/src/test/repl/nativeRepl.test.ts
@@ -1,14 +1,21 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 /* eslint-disable no-unused-expressions */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import * as TypeMoq from 'typemoq';
 import * as sinon from 'sinon';
-import { Disposable } from 'vscode';
+import { Disposable, EventEmitter, NotebookDocument, Uri } from 'vscode';
 import { expect } from 'chai';
 
 import { IInterpreterService } from '../../client/interpreter/contracts';
 import { PythonEnvironment } from '../../client/pythonEnvironments/info';
-import { getNativeRepl, NativeRepl } from '../../client/repl/nativeRepl';
+import * as NativeReplModule from '../../client/repl/nativeRepl';
 import * as persistentState from '../../client/common/persistentState';
+import * as PythonServer from '../../client/repl/pythonServer';
+import * as vscodeWorkspaceApis from '../../client/common/vscodeApis/workspaceApis';
+import * as replController from '../../client/repl/replController';
+import { executeCommand } from '../../client/common/vscodeApis/commandApis';
 
 suite('REPL - Native REPL', () => {
     let interpreterService: TypeMoq.IMock<IInterpreterService>;
@@ -19,38 +26,51 @@ suite('REPL - Native REPL', () => {
     let setReplControllerSpy: sinon.SinonSpy;
     let getWorkspaceStateValueStub: sinon.SinonStub;
     let updateWorkspaceStateValueStub: sinon.SinonStub;
+    let createReplControllerStub: sinon.SinonStub;
+    let mockNotebookController: any;
 
     setup(() => {
+        (NativeReplModule as any).nativeRepl = undefined;
+
+        mockNotebookController = {
+            id: 'mockController',
+            dispose: sinon.stub(),
+            updateNotebookAffinity: sinon.stub(),
+            createNotebookCellExecution: sinon.stub(),
+            variableProvider: null,
+        };
+
         interpreterService = TypeMoq.Mock.ofType<IInterpreterService>();
         interpreterService
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
         disposable = TypeMoq.Mock.ofType<Disposable>();
         disposableArray = [disposable.object];
 
-        setReplDirectoryStub = sinon.stub(NativeRepl.prototype as any, 'setReplDirectory').resolves(); // Stubbing private method
-        // Use a spy instead of a stub for setReplController
-        setReplControllerSpy = sinon.spy(NativeRepl.prototype, 'setReplController');
+        createReplControllerStub = sinon.stub(replController, 'createReplController').returns(mockNotebookController);
+        setReplDirectoryStub = sinon.stub(NativeReplModule.NativeRepl.prototype as any, 'setReplDirectory').resolves();
+        setReplControllerSpy = sinon.spy(NativeReplModule.NativeRepl.prototype, 'setReplController');
         updateWorkspaceStateValueStub = sinon.stub(persistentState, 'updateWorkspaceStateValue').resolves();
     });
 
-    teardown(() => {
+    teardown(async () => {
         disposableArray.forEach((d) => {
             if (d) {
                 d.dispose();
             }
         });
         disposableArray = [];
         sinon.restore();
+        executeCommand('workbench.action.closeActiveEditor');
     });
 
     test('getNativeRepl should call create constructor', async () => {
-        const createMethodStub = sinon.stub(NativeRepl, 'create');
+        const createMethodStub = sinon.stub(NativeReplModule.NativeRepl, 'create');
         interpreterService
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
         const interpreter = await interpreterService.object.getActiveInterpreter();
-        await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
+        await NativeReplModule.getNativeRepl(interpreter as PythonEnvironment, disposableArray);
 
         expect(createMethodStub.calledOnce).to.be.true;
     });
@@ -61,7 +81,7 @@ suite('REPL - Native REPL', () => {
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
         const interpreter = await interpreterService.object.getActiveInterpreter();
-        const nativeRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
+        const nativeRepl = await NativeReplModule.getNativeRepl(interpreter as PythonEnvironment, disposableArray);
 
         nativeRepl.sendToNativeRepl(undefined, false);
 
@@ -74,7 +94,7 @@ suite('REPL - Native REPL', () => {
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
         const interpreter = await interpreterService.object.getActiveInterpreter();
-        const nativeRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
+        const nativeRepl = await NativeReplModule.getNativeRepl(interpreter as PythonEnvironment, disposableArray);
 
         nativeRepl.sendToNativeRepl(undefined, false);
 
@@ -87,34 +107,81 @@ suite('REPL - Native REPL', () => {
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
 
-        await NativeRepl.create(interpreter as PythonEnvironment);
+        await NativeReplModule.NativeRepl.create(interpreter as PythonEnvironment);
 
         expect(setReplDirectoryStub.calledOnce).to.be.true;
         expect(setReplControllerSpy.calledOnce).to.be.true;
-
-        setReplDirectoryStub.restore();
-        setReplControllerSpy.restore();
+        expect(createReplControllerStub.calledOnce).to.be.true;
     });
 
-    test('Closing REPL should dispose of the previous instance and create a new one on reopening', async () => {
-        const disposeSpy = sinon.spy(NativeRepl.prototype, 'dispose');
-        const createStub = sinon.stub(NativeRepl, 'create').callThrough();
-        interpreterService
-            .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
-            .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
-        const interpreter = await interpreterService.object.getActiveInterpreter();
-
-        const firstRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
-
-        firstRepl.dispose();
+    test('watchNotebookClosed should clean up resources when notebook is closed', async () => {
+        const notebookCloseEmitter = new EventEmitter<NotebookDocument>();
+        sinon.stub(vscodeWorkspaceApis, 'onDidCloseNotebookDocument').callsFake((handler) => {
+            const disposable = notebookCloseEmitter.event(handler);
+            return disposable;
+        });
 
-        const secondRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
+        const mockPythonServer = {
+            onCodeExecuted: new EventEmitter<void>().event,
+            execute: sinon.stub().resolves({ status: true, output: 'test output' }),
+            executeSilently: sinon.stub().resolves({ status: true, output: 'test output' }),
+            interrupt: sinon.stub(),
+            input: sinon.stub(),
+            checkValidCommand: sinon.stub().resolves(true),
+            dispose: sinon.stub(),
+        };
+
+        // Track the number of times createPythonServer was called
+        let createPythonServerCallCount = 0;
+        sinon.stub(PythonServer, 'createPythonServer').callsFake(() => {
+            // eslint-disable-next-line no-plusplus
+            createPythonServerCallCount++;
+            return mockPythonServer;
+        });
 
-        expect(disposeSpy.calledOnce).to.be.true;
-        expect(createStub.calledTwice).to.be.true;
-        expect(firstRepl).to.not.equal(secondRepl);
+        const interpreter = await interpreterService.object.getActiveInterpreter();
 
-        disposeSpy.restore();
-        createStub.restore();
+        // Create NativeRepl directly to have more control over its state, go around private constructor.
+        const nativeRepl = new (NativeReplModule.NativeRepl as any)();
+        nativeRepl.interpreter = interpreter as PythonEnvironment;
+        nativeRepl.cwd = '/helloJustMockedCwd/cwd';
+        nativeRepl.pythonServer = mockPythonServer;
+        nativeRepl.replController = mockNotebookController;
+        nativeRepl.disposables = [];
+
+        // Make the singleton point to our instance for testing
+        // Otherwise, it gets mixed with Native Repl from .create from test above.
+        (NativeReplModule as any).nativeRepl = nativeRepl;
+
+        // Reset call count after initial setup
+        createPythonServerCallCount = 0;
+
+        // Set notebookDocument to a mock document
+        const mockReplUri = Uri.parse('untitled:Untitled-999.ipynb?jupyter-notebook');
+        const mockNotebookDocument = ({
+            uri: mockReplUri,
+            toString: () => mockReplUri.toString(),
+        } as unknown) as NotebookDocument;
+
+        nativeRepl.notebookDocument = mockNotebookDocument;
+
+        // Create a mock notebook document for closing event with same URI
+        const closingNotebookDocument = ({
+            uri: mockReplUri,
+            toString: () => mockReplUri.toString(),
+        } as unknown) as NotebookDocument;
+
+        notebookCloseEmitter.fire(closingNotebookDocument);
+        await new Promise((resolve) => setTimeout(resolve, 50));
+
+        expect(
+            updateWorkspaceStateValueStub.calledWith(NativeReplModule.NATIVE_REPL_URI_MEMENTO, undefined),
+            'updateWorkspaceStateValue should be called with NATIVE_REPL_URI_MEMENTO and undefined',
+        ).to.be.true;
+        expect(mockPythonServer.dispose.calledOnce, 'pythonServer.dispose() should be called once').to.be.true;
+        expect(createPythonServerCallCount, 'createPythonServer should be called to create a new server').to.equal(1);
+        expect(nativeRepl.notebookDocument, 'notebookDocument should be undefined after closing').to.be.undefined;
+        expect(nativeRepl.newReplSession, 'newReplSession should be set to true after closing').to.be.true;
+        expect(mockNotebookController.dispose.calledOnce, 'replController.dispose() should be called once').to.be.true;
     });
 });
