From 091e121df020ee7eaafd5e310535245477d40d99 Mon Sep 17 00:00:00 2001
From: Eleanor Boyd <eleanorboyd@microsoft.com>
Date: Fri, 6 Oct 2023 18:34:48 -0700
Subject: [PATCH] fix bug with unittest debug not having args (#22169)

---
 src/client/testing/common/debugLauncher.ts    | 12 +---
 .../testing/testController/common/server.ts   |  2 +
 .../testController/server.unit.test.ts        | 60 ++++++++++++++++++-
 3 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/src/client/testing/common/debugLauncher.ts b/src/client/testing/common/debugLauncher.ts
index 63e2a4543beb..c76557699ff2 100644
--- a/src/client/testing/common/debugLauncher.ts
+++ b/src/client/testing/common/debugLauncher.ts
@@ -205,19 +205,13 @@ export class DebugLauncher implements ITestDebugLauncher {
         }
         launchArgs.request = 'launch';
 
-        // Both types of tests need to have the port for the test result server.
-        if (options.runTestIdsPort) {
-            launchArgs.env = {
-                ...launchArgs.env,
-                RUN_TEST_IDS_PORT: options.runTestIdsPort,
-            };
-        }
-        if (options.testProvider === 'pytest' && pythonTestAdapterRewriteExperiment) {
-            if (options.pytestPort && options.pytestUUID) {
+        if (pythonTestAdapterRewriteExperiment) {
+            if (options.pytestPort && options.pytestUUID && options.runTestIdsPort) {
                 launchArgs.env = {
                     ...launchArgs.env,
                     TEST_PORT: options.pytestPort,
                     TEST_UUID: options.pytestUUID,
+                    RUN_TEST_IDS_PORT: options.runTestIdsPort,
                 };
             } else {
                 throw Error(
diff --git a/src/client/testing/testController/common/server.ts b/src/client/testing/testController/common/server.ts
index 7437a44d6080..50ae1f3f7536 100644
--- a/src/client/testing/testController/common/server.ts
+++ b/src/client/testing/testController/common/server.ts
@@ -211,6 +211,8 @@ export class PythonTestServer implements ITestServer, Disposable {
                     token: options.token,
                     testProvider: UNITTEST_PROVIDER,
                     runTestIdsPort: runTestIdPort,
+                    pytestUUID: uuid.toString(),
+                    pytestPort: this.getPort().toString(),
                 };
                 traceInfo(`Running DEBUG unittest with arguments: ${args}\r\n`);
 
diff --git a/src/test/testing/testController/server.unit.test.ts b/src/test/testing/testController/server.unit.test.ts
index eaf94eca5189..742492b33ba8 100644
--- a/src/test/testing/testController/server.unit.test.ts
+++ b/src/test/testing/testController/server.unit.test.ts
@@ -16,7 +16,7 @@ import {
     Output,
 } from '../../../client/common/process/types';
 import { PythonTestServer } from '../../../client/testing/testController/common/server';
-import { ITestDebugLauncher } from '../../../client/testing/common/types';
+import { ITestDebugLauncher, LaunchOptions } from '../../../client/testing/common/types';
 import { Deferred, createDeferred } from '../../../client/common/utils/async';
 import { MockChildProcess } from '../../mocks/mockChildProcess';
 import {
@@ -240,6 +240,64 @@ suite('Python Test Server, Send command etc', () => {
             assert(false, errorMessage);
         }
     });
+    test('sendCommand should add right extra variables to command during debug', async () => {
+        const deferred2 = createDeferred();
+        const RUN_TEST_IDS_PORT_CONST = '5678';
+        const error = false;
+        const errorMessage = '';
+        const debugLauncherMock = typeMoq.Mock.ofType<ITestDebugLauncher>();
+        let actualLaunchOptions: LaunchOptions = {} as LaunchOptions;
+        const deferred4 = createDeferred();
+        debugLauncherMock
+            .setup((x) => x.launchDebugger(typeMoq.It.isAny(), typeMoq.It.isAny()))
+            .returns((options, _) => {
+                actualLaunchOptions = options;
+                deferred4.resolve();
+                return Promise.resolve();
+            });
+        execService
+            .setup((x) => x.execObservable(typeMoq.It.isAny(), typeMoq.It.isAny()))
+            .returns(() => typeMoq.Mock.ofType<ObservableExecutionResult<string>>().object);
+        const execFactory = typeMoq.Mock.ofType<IPythonExecutionFactory>();
+        execFactory
+            .setup((x) => x.createActivatedEnvironment(typeMoq.It.isAny()))
+            .returns(() => {
+                deferred2.resolve();
+                return Promise.resolve(execService.object);
+            });
+        server = new PythonTestServer(execFactory.object, debugLauncherMock.object);
+        sinon.stub(server, 'getPort').returns(12345);
+        // const portServer = server.getPort();
+        await server.serverReady();
+        const options = {
+            command: { script: 'myscript', args: ['-foo', 'foo'] },
+            workspaceFolder: Uri.file('/foo/bar'),
+            cwd: '/foo/bar',
+            uuid: FAKE_UUID,
+            debugBool: true,
+        };
+        try {
+            server.sendCommand(options, {}, RUN_TEST_IDS_PORT_CONST);
+        } catch (e) {
+            assert(false, `Error sending command, ${e}`);
+        }
+        // add in await and trigger
+        await deferred2.promise;
+        await deferred4.promise;
+        mockProc.trigger('close');
+
+        assert.notDeepEqual(actualLaunchOptions, {}, 'launch options should be set');
+        assert.strictEqual(actualLaunchOptions.cwd, '/foo/bar');
+        assert.strictEqual(actualLaunchOptions.testProvider, 'unittest');
+        assert.strictEqual(actualLaunchOptions.pytestPort, '12345');
+        assert.strictEqual(actualLaunchOptions.pytestUUID, 'fake-uuid');
+        assert.strictEqual(actualLaunchOptions.runTestIdsPort, '5678');
+
+        debugLauncherMock.verify((x) => x.launchDebugger(typeMoq.It.isAny(), typeMoq.It.isAny()), typeMoq.Times.once());
+        if (error) {
+            assert(false, errorMessage);
+        }
+    });
 
     test('sendCommand should write to an output channel if it is provided as an option', async () => {
         const output2: string[] = [];