mirror of https://github.com/databricks/cli.git
108 lines
2.5 KiB
Go
108 lines
2.5 KiB
Go
package pythontest
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/databricks/cli/internal/testutil"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type VenvOpts struct {
|
|
// input
|
|
PythonVersion string
|
|
skipVersionCheck bool
|
|
|
|
// 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-")
|
|
}
|
|
|
|
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)
|
|
|
|
newPath := fmt.Sprintf("%s%c%s", opts.BinPath, os.PathListSeparator, os.Getenv("PATH"))
|
|
t.Setenv("PATH", newPath)
|
|
}
|