better sanity check, more test

Author: anthonykim1

src/client/repl/nativeRepl.ts

@@ -169,15 +169,19 @@ export class NativeRepl implements Disposable {
             wsMementoUri,
             preserveFocus,
         );
+        if (notebookEditor) {
+            this.notebookDocument = notebookEditor.notebook;
+            updateWorkspaceStateValue<string | undefined>(
+                NATIVE_REPL_URI_MEMENTO,
+                this.notebookDocument.uri.toString(),
+            );
 
-        this.notebookDocument = notebookEditor.notebook;
-        updateWorkspaceStateValue<string | undefined>(NATIVE_REPL_URI_MEMENTO, this.notebookDocument.uri.toString());
-
-        if (this.notebookDocument) {
-            this.replController.updateNotebookAffinity(this.notebookDocument, NotebookControllerAffinity.Default);
-            await selectNotebookKernel(notebookEditor, this.replController.id, PVSC_EXTENSION_ID);
-            if (code) {
-                await executeNotebookCell(notebookEditor, code);
+            if (this.notebookDocument) {
+                this.replController.updateNotebookAffinity(this.notebookDocument, NotebookControllerAffinity.Default);
+                await selectNotebookKernel(notebookEditor, this.replController.id, PVSC_EXTENSION_ID);
+                if (code) {
+                    await executeNotebookCell(notebookEditor, code);
+                }
             }
         }
     }

src/client/repl/replCommandHandler.ts

@@ -12,7 +12,7 @@ import {
     workspace,
     Uri,
 } from 'vscode';
-import { getExistingReplViewColumn } from './replUtils';
+import { getExistingReplViewColumn, getTabNameForUri } from './replUtils';
 import { PVSC_EXTENSION_ID } from '../common/constants';
 
 /**
@@ -23,7 +23,7 @@ export async function openInteractiveREPL(
     notebookDocument: NotebookDocument | undefined,
     mementoValue: Uri | undefined,
     preserveFocus: boolean = true,
-): Promise<NotebookEditor> {
+): Promise<NotebookEditor | undefined> {
     let viewColumn = ViewColumn.Beside;
     if (mementoValue) {
         if (!notebookDocument) {
@@ -38,13 +38,22 @@ export async function openInteractiveREPL(
         // became outdated (untitled.ipynb created without Python extension knowing, effectively taking over original Python REPL's URI)
         notebookDocument = await workspace.openNotebookDocument('jupyter-notebook');
     }
-    const editor = window.showNotebookDocument(notebookDocument!, {
+    const editor = await window.showNotebookDocument(notebookDocument!, {
         viewColumn,
         asRepl: 'Python REPL',
         preserveFocus,
     });
-    // sanity check that we opened a Native REPL from showNotebookDocument.
-    // if not true, set notebook = undefined.
+
+    // Sanity check that we opened a Native REPL from showNotebookDocument.
+    if (
+        !editor ||
+        !editor.notebook ||
+        !editor.notebook.uri ||
+        getTabNameForUri(editor.notebook.uri) !== 'Python REPL'
+    ) {
+        return undefined;
+    }
+
     await commands.executeCommand('notebook.selectKernel', {
         editor,
         id: notebookController.id,

src/test/repl/nativeRepl.test.ts

@@ -18,6 +18,7 @@ suite('REPL - Native REPL', () => {
     let setReplDirectoryStub: sinon.SinonStub;
     let setReplControllerSpy: sinon.SinonSpy;
     let getWorkspaceStateValueStub: sinon.SinonStub;
+    let updateWorkspaceStateValueStub: sinon.SinonStub;
 
     setup(() => {
         interpreterService = TypeMoq.Mock.ofType<IInterpreterService>();
@@ -30,7 +31,7 @@ suite('REPL - Native REPL', () => {
         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');
-        getWorkspaceStateValueStub = sinon.stub(persistentState, 'getWorkspaceStateValue').returns(undefined);
+        updateWorkspaceStateValueStub = sinon.stub(persistentState, 'updateWorkspaceStateValue').resolves();
     });
 
     teardown(() => {
@@ -55,6 +56,7 @@ suite('REPL - Native REPL', () => {
     });
 
     test('sendToNativeRepl should look for memento URI if notebook document is undefined', async () => {
+        getWorkspaceStateValueStub = sinon.stub(persistentState, 'getWorkspaceStateValue').returns(undefined);
         interpreterService
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
@@ -66,6 +68,19 @@ suite('REPL - Native REPL', () => {
         expect(getWorkspaceStateValueStub.calledOnce).to.be.true;
     });
 
+    test('sendToNativeRepl should call updateWorkspaceStateValue', async () => {
+        getWorkspaceStateValueStub = sinon.stub(persistentState, 'getWorkspaceStateValue').returns('myNameIsMemento');
+        interpreterService
+            .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);
+
+        nativeRepl.sendToNativeRepl(undefined, false);
+
+        expect(updateWorkspaceStateValueStub.calledOnce).to.be.true;
+    });
+
     test('create should call setReplDirectory, setReplController', async () => {
         const interpreter = await interpreterService.object.getActiveInterpreter();
         interpreterService