From 3c3c172622b22c7868cf1b82bb02caeb748e2875 Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Wed, 18 Dec 2024 20:20:30 +0100
Subject: [PATCH] Fix finding Python within virtualenv on Windows; new
 pythontest package

---
 .github/workflows/push.yml                |  3 +++
 internal/testutil/cmd.go                  | 26 ++++++++++++++++++
 internal/testutil/env.go                  | 22 ++++++++++++++++
 libs/python/detect.go                     | 17 ++++++++++++
 libs/python/pythontest/pythontest.go      | 32 +++++++++++++++++++++++
 libs/python/pythontest/pythontest_test.go | 16 ++++++++++++
 6 files changed, 116 insertions(+)
 create mode 100644 internal/testutil/cmd.go
 create mode 100644 libs/python/pythontest/pythontest.go
 create mode 100644 libs/python/pythontest/pythontest_test.go

diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml
index ea561805b..e5bc6b192 100644
--- a/.github/workflows/push.yml
+++ b/.github/workflows/push.yml
@@ -40,6 +40,9 @@ jobs:
         with:
           python-version: '3.9'
 
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
       - name: Set go env
         run: |
           echo "GOPATH=$(go env GOPATH)" >> $GITHUB_ENV
diff --git a/internal/testutil/cmd.go b/internal/testutil/cmd.go
new file mode 100644
index 000000000..4045c606c
--- /dev/null
+++ b/internal/testutil/cmd.go
@@ -0,0 +1,26 @@
+package testutil
+
+import (
+	"bytes"
+	"os"
+	"os/exec"
+
+	"github.com/stretchr/testify/require"
+)
+
+func RunCommand(t TestingT, name string, args ...string) {
+	cmd := exec.Command(name, args...)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	require.NoError(t, cmd.Run())
+}
+
+func CaptureCommandOutput(t TestingT, name string, args ...string) string {
+	cmd := exec.Command(name, args...)
+	var stdout bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = os.Stderr
+	err := cmd.Run()
+	require.NoError(t, err)
+	return stdout.String()
+}
diff --git a/internal/testutil/env.go b/internal/testutil/env.go
index 10557c4e6..0ac3fbea8 100644
--- a/internal/testutil/env.go
+++ b/internal/testutil/env.go
@@ -61,3 +61,25 @@ func Chdir(t TestingT, dir string) string {
 
 	return wd
 }
+
+func InsertPathEntry(t TestingT, path string) {
+	var separator string
+	if runtime.GOOS == "windows" {
+		separator = ";"
+	} else {
+		separator = ":"
+	}
+
+	t.Setenv("PATH", path+separator+os.Getenv("PATH"))
+}
+
+func InsertVirtualenvInPath(t TestingT, venvPath string) {
+	if runtime.GOOS == "windows" {
+		// https://github.com/pypa/virtualenv/commit/993ba1316a83b760370f5a3872b3f5ef4dd904c1
+		venvPath = filepath.Join(venvPath, "Scripts")
+	} else {
+		venvPath = filepath.Join(venvPath, "bin")
+	}
+
+	InsertPathEntry(t, venvPath)
+}
diff --git a/libs/python/detect.go b/libs/python/detect.go
index 8fcc7cd9c..247548a71 100644
--- a/libs/python/detect.go
+++ b/libs/python/detect.go
@@ -25,15 +25,32 @@ func DetectExecutable(ctx context.Context) (string, error) {
 	// the parent directory tree.
 	//
 	// See https://github.com/pyenv/pyenv#understanding-python-version-selection
+
+	// On Windows when virtualenv is created, the <env>/Scripts directory
+	// contains python.exe but no python3.exe. However, system python does have python3 entry
+	// and it is also added to PATH, so it is found first.
+	if runtime.GOOS == "windows" {
+		out, err := exec.LookPath("python.exe")
+		if err == nil && out != "" {
+			return out, nil
+		}
+		if err != nil && !errors.Is(err, exec.ErrNotFound) {
+			return "", err
+		}
+	}
+
 	out, err := exec.LookPath("python3")
+
 	// most of the OS'es have python3 in $PATH, but for those which don't,
 	// we perform the latest version lookup
 	if err != nil && !errors.Is(err, exec.ErrNotFound) {
 		return "", err
 	}
+
 	if out != "" {
 		return out, nil
 	}
+
 	// otherwise, detect all interpreters and pick the least that satisfies
 	// minimal version requirements
 	all, err := DetectInterpreters(ctx)
diff --git a/libs/python/pythontest/pythontest.go b/libs/python/pythontest/pythontest.go
new file mode 100644
index 000000000..5afea6366
--- /dev/null
+++ b/libs/python/pythontest/pythontest.go
@@ -0,0 +1,32 @@
+package pythontest
+
+import (
+	"context"
+	"path/filepath"
+	"strings"
+
+	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/python"
+	"github.com/stretchr/testify/require"
+)
+
+func RequirePythonVENV(t testutil.TestingT, ctx context.Context, pythonVersion string, checkVersion bool) string {
+	tmpDir := t.TempDir()
+	testutil.Chdir(t, tmpDir)
+
+	venvName := testutil.RandomName("test-venv-")
+	testutil.RunCommand(t, "uv", "venv", venvName, "--python", pythonVersion, "--seed")
+	testutil.InsertVirtualenvInPath(t, filepath.Join(tmpDir, venvName))
+
+	pythonExe, err := python.DetectExecutable(ctx)
+	require.NoError(t, err)
+	require.Contains(t, pythonExe, venvName)
+
+	if checkVersion {
+		actualVersion := testutil.CaptureCommandOutput(t, pythonExe, "--version")
+		expectVersion := "Python " + pythonVersion
+		require.True(t, strings.HasPrefix(actualVersion, expectVersion), "Running %s --version: Expected %v, got %v", pythonExe, expectVersion, actualVersion)
+	}
+
+	return tmpDir
+}
diff --git a/libs/python/pythontest/pythontest_test.go b/libs/python/pythontest/pythontest_test.go
new file mode 100644
index 000000000..3aba780c5
--- /dev/null
+++ b/libs/python/pythontest/pythontest_test.go
@@ -0,0 +1,16 @@
+package pythontest
+
+import (
+	"context"
+	"testing"
+)
+
+func TestVenv(t *testing.T) {
+	// Test at least two version to ensure we capture a case where venv version does not match system one
+	for _, pythonVersion := range []string{"3.11", "3.12"} {
+		t.Run(pythonVersion, func(t *testing.T) {
+			ctx := context.Background()
+			RequirePythonVENV(t, ctx, pythonVersion, true)
+		})
+	}
+}