From b4273f96248c2af4d9f7d700f367b24dbc51369f Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Tue, 4 Mar 2025 15:37:58 +0100
Subject: [PATCH] fix re-patching and add test

---
 libs/patchwheel/patch.go      |  9 +++------
 libs/patchwheel/patch_test.go | 17 ++++++++++++-----
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/libs/patchwheel/patch.go b/libs/patchwheel/patch.go
index 59e8abf83..1ce59cd2d 100644
--- a/libs/patchwheel/patch.go
+++ b/libs/patchwheel/patch.go
@@ -185,12 +185,9 @@ func PatchWheel(ctx context.Context, path, outputDir string) (string, error) {
 	metadataSize := len(newMetadata)
 
 	// Compute the new dist-info directory prefix.
-	newDistInfoPrefix := ""
-	if idx := strings.LastIndex(oldDistInfoPrefix, "-"); idx != -1 {
-		base := oldDistInfoPrefix[:idx]
-		newDistInfoPrefix = base + "-" + newVersion + ".dist-info/"
-	} else {
-		return "", fmt.Errorf("unexpected dist-info directory format: %s", oldDistInfoPrefix)
+	newDistInfoPrefix := strings.Replace(oldDistInfoPrefix, wheelInfo.Version, newVersion, 1)
+	if newDistInfoPrefix == oldDistInfoPrefix {
+		return "", fmt.Errorf("unexpected dist-info directory format: %s (version=%s)", oldDistInfoPrefix, wheelInfo.Version)
 	}
 
 	recordReader, err := recordFile.Open()
diff --git a/libs/patchwheel/patch_test.go b/libs/patchwheel/patch_test.go
index b8c0cd351..da1febabe 100644
--- a/libs/patchwheel/patch_test.go
+++ b/libs/patchwheel/patch_test.go
@@ -8,6 +8,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -120,11 +121,9 @@ func TestPatchWheel(t *testing.T) {
 			distDir := filepath.Join(tempDir, "dist")
 			origWheel := getWheel(t, distDir)
 
-			// First patch
 			patchedWheel, err := PatchWheel(context.Background(), origWheel, distDir)
 			require.NoError(t, err)
 
-			// Get file info of the patched wheel
 			patchedInfo, err := os.Stat(patchedWheel)
 			require.NoError(t, err)
 			patchedTime := patchedInfo.ModTime()
@@ -141,11 +140,19 @@ func TestPatchWheel(t *testing.T) {
 			require.Equal(t, patchedTime, patchedInfo2.ModTime(), "File was recreated when it shouldn't have been")
 
 			runCmd(t, tempDir, "uv", "pip", "install", "-q", patchedWheel)
-
-			// Verify the installed version matches what we expect
 			verifyVersion(t, tempDir, patchedWheel)
 
-			// TODO: install one more patched wheel (add an option to PatchWheel to add extra to timestamp)
+			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)
+
+			runCmd(t, tempDir, "uv", "pip", "install", "-q", patchedWheel3)
+			verifyVersion(t, tempDir, patchedWheel3)
 		})
 	}
 }