From dcfeeacea2a1610f43c3d73acf80596f2d18a9ac Mon Sep 17 00:00:00 2001 From: Denis Bilenko Date: Thu, 19 Dec 2024 10:06:17 +0100 Subject: [PATCH] refactor: split venv creation into a separate function this new function does not change chdir or sets environment --- internal/testutil/cmd.go | 26 ----- libs/python/pythontest/pythontest.go | 112 +++++++++++++++++++--- libs/python/pythontest/pythontest_test.go | 19 +++- 3 files changed, 114 insertions(+), 43 deletions(-) delete mode 100644 internal/testutil/cmd.go diff --git a/internal/testutil/cmd.go b/internal/testutil/cmd.go deleted file mode 100644 index 4045c606c..000000000 --- a/internal/testutil/cmd.go +++ /dev/null @@ -1,26 +0,0 @@ -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/libs/python/pythontest/pythontest.go b/libs/python/pythontest/pythontest.go index 5afea6366..7e5e2eb28 100644 --- a/libs/python/pythontest/pythontest.go +++ b/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 + + // 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 !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 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.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 + require.Equal(t, filepath.Dir(pythonExe), filepath.Dir(opts.PythonExe)) } diff --git a/libs/python/pythontest/pythontest_test.go b/libs/python/pythontest/pythontest_test.go index 3aba780c5..d8eeca64b 100644 --- a/libs/python/pythontest/pythontest_test.go +++ b/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{})) +}