2 files changed, 72 insertions, 15 deletions
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index 4e6be9b3..71b774d9 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -4,7 +4,6 @@ const exec = util.promisify(require("child_process").exec);
 const { spawn } = require("child_process");
 import * as path from "path";
 import * as fs from "fs";
-import rebuild from "@electron/rebuild";
 import { getContinueServerUrl } from "../bridge";
 import fetch from "node-fetch";
 import * as vscode from "vscode";
@@ -60,7 +59,7 @@ async function runCommand(cmd: string): Promise<[string, string | undefined]> {
   return [stdout, stderr];
 }
 
-async function getPythonPipCommands() {
+export async function getPythonPipCommands() {
   var [stdout, stderr] = await runCommand("python3 --version");
   let pythonCmd = "python3";
   if (stderr) {
@@ -111,10 +110,8 @@ async function getPythonPipCommands() {
       );
       throw new Error("Python3.8 or greater is not installed.");
     }
-  } else {
-    pythonCmd = `python${major}.${minor}`;
-    pipCmd = `pip${major}.${minor}`;
   }
+
   return [pythonCmd, pipCmd];
 }
 
@@ -368,16 +365,6 @@ export async function startContinuePythonServer() {
   });
 }
 
-async function installNodeModules() {
-  console.log("Rebuilding node-pty for Continue extension...");
-  await rebuild({
-    buildPath: getExtensionUri().fsPath, // Folder containing node_modules
-    electronVersion: "19.1.8",
-    onlyModules: ["node-pty"],
-  });
-  console.log("Successfully rebuilt node-pty");
-}
-
 export function isPythonEnvSetup(): boolean {
   let pathToEnvCfg = getExtensionUri().fsPath + "/scripts/env/pyvenv.cfg";
   return fs.existsSync(path.join(pathToEnvCfg));
diff --git a/extension/src/activation/test/environmentSetup.test.ts b/extension/src/activation/test/environmentSetup.test.ts
new file mode 100644
index 00000000..ca487416
--- /dev/null
+++ b/extension/src/activation/test/environmentSetup.test.ts
@@ -0,0 +1,70 @@
+const child_process = require("child_process");
+import { platform } from "os";
+import { getPythonPipCommands } from "../environmentSetup";
+
+jest.mock("os");
+jest.mock("child_process");
+
+function mockPythonVersionMappings(mappings: { [pythonCmd: string]: string }) {
+  (child_process.exec as jest.Mock).mockImplementation(
+    (command: string, options: any) => {
+      const pythonCmd = command.split(" ")[0];
+      if (pythonCmd in mappings) {
+        return Promise.resolve([mappings[pythonCmd], ""]);
+      } else {
+        return Promise.resolve(["", stubStderr]);
+      }
+    }
+  );
+}
+
+const stubStderr =
+  "This is a stub stderr, but will be checked only for existence.";
+describe("getPythonPipCommands", () => {
+  describe("on Windows", () => {
+    it("should return the correct Python and Pip commands", async () => {
+      (platform as jest.Mock).mockReturnValue("win32");
+      mockPythonVersionMappings({
+        python: "Python 3.8.0",
+      });
+
+      const [pythonCmd, pipCmd] = await getPythonPipCommands();
+
+      expect(pythonCmd).toBe("python");
+      expect(pipCmd).toBe("pip");
+
+      jest.restoreAllMocks();
+    });
+    describe("on MacOS", () => {
+      (platform as jest.Mock).mockReturnValue("darwin");
+      it("should check through all python versions after finding 3.7", async () => {
+        mockPythonVersionMappings({
+          python: "",
+          python3: "Python 3.7.0",
+          "python3.11": "Python 3.11.0",
+        });
+
+        const [pythonCmd, pipCmd] = await getPythonPipCommands();
+
+        expect(pythonCmd).toBe("python3.11");
+        expect(pipCmd).toBe("pip3.11");
+
+        jest.restoreAllMocks();
+      });
+
+      it("should use python3 if that maps to valid version", async () => {
+        mockPythonVersionMappings({
+          python: "",
+          python3: "Python 3.8.0",
+        });
+
+        const [pythonCmd, pipCmd] = await getPythonPipCommands();
+
+        expect(pythonCmd).toBe("python3");
+        expect(pipCmd).toBe("pip3");
+
+        jest.restoreAllMocks();
+      });
+    });
+  });
+});