refactor: split venv creation into a separate function

this new function does not change chdir or sets environment

libs/python/pythontest/pythontest.go

@@ -2,31 +2,113 @@ package pythontest
 
 import (
 	"context"
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
 	"path/filepath"
+	"runtime"
 	"strings"
+	"testing"
 
 	"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)
+type VenvOpts struct {
+	// input
+	PythonVersion    string
+	skipVersionCheck bool
 
-	venvName := testutil.RandomName("test-venv-")
-	testutil.RunCommand(t, "uv", "venv", venvName, "--python", pythonVersion, "--seed")
-	testutil.InsertVirtualenvInPath(t, filepath.Join(tmpDir, venvName))
+	// input/output
+	Dir  string
+	Name string
 
-	pythonExe, err := python.DetectExecutable(ctx)
-	require.NoError(t, err)
-	require.Contains(t, pythonExe, venvName)
+	// output:
+	// Absolute path to venv
+	EnvPath string
+
+	// Absolute path to venv/bin or venv/Scripts, depending on OS
+	BinPath string
+
+	// Absolute path to python binary
+	PythonExe string
+}
+
+func CreatePythonEnv(opts *VenvOpts) error {
+	if opts == nil || opts.PythonVersion == "" {
+		return errors.New("PythonVersion must be provided")
+	}
+	if opts.Name == "" {
+		opts.Name = testutil.RandomName("test-venv-")
+	}
+	if opts.Dir != "" {
+		opts.Dir = "."
+	}
+
+	cmd := exec.Command("uv", "venv", opts.Name, "--python", opts.PythonVersion, "--seed", "-q")
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	cmd.Dir = opts.Dir
+	err := cmd.Run()
+	if err != nil {
+		return err
+	}
+
+	opts.EnvPath, err = filepath.Abs(filepath.Join(opts.Dir, opts.Name))
+	if err != nil {
+		return err
+	}
+
+	_, err = os.Stat(opts.EnvPath)
+	if err != nil {
+		return fmt.Errorf("cannot stat EnvPath %s: %s", opts.EnvPath, err)
+	}
+
+	if runtime.GOOS == "windows" {
+		// https://github.com/pypa/virtualenv/commit/993ba1316a83b760370f5a3872b3f5ef4dd904c1
+		opts.BinPath = filepath.Join(opts.EnvPath, "Scripts")
+		opts.PythonExe = filepath.Join(opts.BinPath, "python.exe")
+	} else {
+		opts.BinPath = filepath.Join(opts.EnvPath, "bin")
+		opts.PythonExe = filepath.Join(opts.BinPath, "python3")
+	}
+
+	_, err = os.Stat(opts.BinPath)
+	if err != nil {
+		return fmt.Errorf("cannot stat BinPath %s: %s", opts.BinPath, err)
+	}
+
+	_, err = os.Stat(opts.PythonExe)
+	if err != nil {
+		return fmt.Errorf("cannot stat PythonExe %s: %s", opts.PythonExe, err)
+	}
 
-	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)
+	if !opts.skipVersionCheck {
+		cmd := exec.Command(opts.PythonExe, "--version")
+		out, err := cmd.CombinedOutput()
+		if err != nil {
+			return fmt.Errorf("Failed to run %s --version: %s", opts.PythonExe, err)
+		}
+		outString := string(out)
+		expectVersion := "Python " + opts.PythonVersion
+		if !strings.HasPrefix(outString, expectVersion) {
+			return fmt.Errorf("Unexpected output from %s --version: %v (expected %v)", opts.PythonExe, outString, expectVersion)
+		}
 	}
 
-	return tmpDir
+	return nil
+}
+
+func RequireActivatedPythonEnv(t *testing.T, ctx context.Context, opts *VenvOpts) {
+	err := CreatePythonEnv(opts)
+	require.NoError(t, err)
+	require.DirExists(t, opts.BinPath)
+
+	testutil.InsertPathEntry(t, opts.BinPath)
+
+	pythonExe, err := python.DetectExecutable(ctx)
+	require.NoError(t, err)
+	require.Equal(t, filepath.Dir(pythonExe), filepath.Dir(opts.PythonExe))
 }

libs/python/pythontest/pythontest_test.go

@@ -3,14 +3,29 @@ package pythontest
 import (
 	"context"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
-func TestVenv(t *testing.T) {
+func TestVenvSuccess(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)
+			opts := VenvOpts{PythonVersion: pythonVersion}
+			RequireActivatedPythonEnv(t, ctx, &opts)
+			require.DirExists(t, opts.EnvPath)
+			require.DirExists(t, opts.BinPath)
+			require.FileExists(t, opts.PythonExe)
 		})
 	}
 }
+
+func TestWrongVersion(t *testing.T) {
+	require.Error(t, CreatePythonEnv(&VenvOpts{PythonVersion: "4.0"}))
+}
+
+func TestMissingVersion(t *testing.T) {
+	require.Error(t, CreatePythonEnv(nil))
+	require.Error(t, CreatePythonEnv(&VenvOpts{}))
+}