From cc59dec5f5e9cef168cfb202892d894f83c045a5 Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Tue, 4 Mar 2025 16:18:34 +0100
Subject: [PATCH] parallel tests and short/non-short mode

---
 libs/patchwheel/patch_test.go | 31 ++++++++++++++++++++++++-------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/libs/patchwheel/patch_test.go b/libs/patchwheel/patch_test.go
index 5aeae890d..5c6873e47 100644
--- a/libs/patchwheel/patch_test.go
+++ b/libs/patchwheel/patch_test.go
@@ -15,7 +15,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-var scriptsDir = getPythonScriptsDir()
+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" {
@@ -24,6 +28,23 @@ func getPythonScriptsDir() string {
 	return "bin"
 }
 
+func getPythonVersions() []string {
+	if testing.Short() {
+		return []string{
+			"python3.9",
+			"python3.12",
+		}
+	} else {
+		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)
@@ -116,9 +137,10 @@ func getWheel(t *testing.T, dir string) string {
 }
 
 func TestPatchWheel(t *testing.T) {
-	pythonVersions := []string{"python3.9", "python3.10", "python3.11", "python3.12"}
+	pythonVersions := getPythonVersions()
 	for _, py := range pythonVersions {
 		t.Run(py, func(t *testing.T) {
+			t.Parallel()
 			tempDir := t.TempDir()
 
 			projFiles := minimalPythonProject()
@@ -165,11 +187,6 @@ func TestPatchWheel(t *testing.T) {
 	}
 }
 
-var (
-	prebuiltWheel = "testdata/my_test_code-0.0.1-py3-none-any.whl"
-	emptyZip      = "testdata/empty.zip"
-)
-
 func TestPrebuilt(t *testing.T) {
 	tempDir := t.TempDir()
 	ctx := context.Background()