package patchwheel import ( "bytes" "context" "os" "os/exec" "path/filepath" "runtime" "strings" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) var ( scriptsDir = getPythonScriptsDir() prebuiltWheel = "testdata/my_test_code-0.0.1-py3-none-any.whl" emptyZip = "testdata/empty.zip" ) func getPythonScriptsDir() string { if runtime.GOOS == "windows" { return "Scripts" } return "bin" } func getPythonVersions() []string { return []string{ "python3.9", "python3.10", "python3.11", "python3.12", "python3.13", } } func verifyVersion(t *testing.T, tempDir, wheelPath string) { wheelInfo, err := ParseWheelFilename(wheelPath) require.NoError(t, err) expectedVersion := wheelInfo.Version pyExec := filepath.Join(tempDir, ".venv", scriptsDir, "python") cmdOut := captureOutput(t, tempDir, pyExec, "-c", "import myproj; myproj.print_version()") actualVersion := strings.TrimSpace(cmdOut) t.Logf("Verified installed version: %s", actualVersion) assert.Equal(t, expectedVersion, actualVersion, "Installed version doesn't match expected version from wheel filename") } func minimalPythonProject() map[string]string { return map[string]string{ "pyproject.toml": `[project] name = "myproj" version = "0.1.0" [build-system] requires = ["setuptools>=61.0.0", "wheel"] build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] where = ["src"] `, "src/myproj/__init__.py": ` def hello(): return "Hello, world!" def print_version(): from importlib.metadata import version print(version("myproj")) `, } } func writeProjectFiles(t *testing.T, baseDir string, files map[string]string) { for path, content := range files { fullPath := filepath.Join(baseDir, path) if err := os.MkdirAll(filepath.Dir(fullPath), 0o755); err != nil { t.Fatalf("Failed to create directory %s: %v", filepath.Dir(fullPath), err) } if err := os.WriteFile(fullPath, []byte(content), 0o644); err != nil { t.Fatalf("Failed to write file %s: %v", fullPath, err) } } } func runCmd(t *testing.T, dir, name string, args ...string) { out := captureOutput(t, dir, name, args...) if len(out) > 0 { t.Errorf("Output from %s %s:\n%s", name, args, out) } } func captureOutput(t *testing.T, dir, name string, args ...string) string { cmd := exec.Command(name, args...) cmd.Dir = dir var out bytes.Buffer cmd.Stdout = &out cmd.Stderr = &out err := cmd.Run() if err != nil { t.Logf("Command failed: %s %s", name, strings.Join(args, " ")) t.Logf("Output:\n%s", out.String()) t.Fatal(err) } return out.String() } func getWheel(t *testing.T, dir string) string { pattern := filepath.Join(dir, "*.whl") matches, err := filepath.Glob(pattern) if err != nil { t.Fatalf("Error matching pattern %s: %v", pattern, err) } if len(matches) == 0 { t.Fatalf("No files found matching %s", pattern) return "" } if len(matches) != 1 { t.Fatalf("Too many matches %s: %v", pattern, matches) return "" } return matches[0] } func TestPatchWheel(t *testing.T) { pythonVersions := getPythonVersions() for _, py := range pythonVersions { t.Run(py, func(t *testing.T) { t.Parallel() tempDir := t.TempDir() projFiles := minimalPythonProject() writeProjectFiles(t, tempDir, projFiles) runCmd(t, tempDir, "uv", "venv", "-q", "--python", py) runCmd(t, tempDir, "uv", "build", "-q", "--wheel") distDir := filepath.Join(tempDir, "dist") origWheel := getWheel(t, distDir) patchedWheel, err := PatchWheel(context.Background(), origWheel, distDir) require.NoError(t, err) patchedInfo, err := os.Stat(patchedWheel) require.NoError(t, err) patchedTime := patchedInfo.ModTime() // Test idempotency - patching the same wheel again should produce the same result // and should not recreate the file (file modification time should remain the same) patchedWheel2, err := PatchWheel(context.Background(), origWheel, distDir) require.NoError(t, err) require.Equal(t, patchedWheel, patchedWheel2, "PatchWheel is not idempotent") patchedInfo2, err := os.Stat(patchedWheel2) require.NoError(t, err) require.Equal(t, patchedTime, patchedInfo2.ModTime(), "File was recreated when it shouldn't have been") runCmd(t, tempDir, "uv", "pip", "install", "-q", patchedWheel) verifyVersion(t, tempDir, patchedWheel) newTime := patchedInfo.ModTime().Add(10 * time.Millisecond) err = os.Chtimes(origWheel, newTime, newTime) require.NoError(t, err) patchedWheel3, err := PatchWheel(context.Background(), origWheel, distDir) require.NoError(t, err) require.Greater(t, patchedWheel3, patchedWheel) // Now use regular pip to re-install the wheel. First install pip. runCmd(t, tempDir, "uv", "pip", "install", "-q", "pip") pippath := filepath.Join(".venv", getPythonScriptsDir(), "pip") runCmd(t, tempDir, pippath, "install", "-q", patchedWheel3) verifyVersion(t, tempDir, patchedWheel3) }) } } func errPatchWheel(t *testing.T, name, out string) { ctx := context.Background() outname, err := PatchWheel(ctx, name, out) assert.Error(t, err, "PatchWheel(%s, %s) expected to error", name, out) assert.Empty(t, outname) } func TestError(t *testing.T) { // empty name and dir errPatchWheel(t, "", "") // empty name errPatchWheel(t, "", ".") // file not found errPatchWheel(t, "not-found.txt", ".") // output directory not found errPatchWheel(t, prebuiltWheel, "not-found/a/b/c") } func TestEmptyZip(t *testing.T) { tempDir := t.TempDir() errPatchWheel(t, emptyZip, tempDir) } func TestNonZip(t *testing.T) { tempDir := t.TempDir() _, err := os.Stat("patch.go") require.NoError(t, err, "file must exist for this test") errPatchWheel(t, "patch.go", tempDir) }