Add integration test for 3.13 base images (#2485)

* Add integration test for 3.13 base images

* Add fallback to binary base64 URL decode

* Python 3.13 adds some mime types that dataurl
cannot recognise.
* To support this, we detect if an error occurs
during decoding and attempt to fix the URL so it
is a generic binary decode.

Author: 8W9aG

pkg/util/files/files.go

@@ -16,6 +16,10 @@ import (
 	"github.com/replicate/cog/pkg/util/mime"
 )
 
+var (
+	ErrorFailedToSplitDataURL = errors.New("Failed to split data URL into 2 parts")
+)
+
 func Exists(path string) (bool, error) {
 	if _, err := os.Stat(path); err == nil {
 		return true, nil
@@ -96,7 +100,17 @@ func WriteDataURLToFile(url string, destination string) (string, error) {
 	}
 	dataurlObj, err := dataurl.DecodeString(url)
 	if err != nil {
-		return "", fmt.Errorf("Failed to decode data URL: %w", err)
+		// Attempt to fallback to binary base64 file decode.
+		parts := strings.SplitN(url, ",", 2)
+		if len(parts) != 2 {
+			return "", ErrorFailedToSplitDataURL
+		}
+		base64Data := parts[1]
+		url = "data:;base64," + base64Data
+		dataurlObj, err = dataurl.DecodeString(url)
+		if err != nil {
+			return "", fmt.Errorf("Failed to decode data URL: %w", err)
+		}
 	}
 	output := dataurlObj.Data
 

pkg/util/files/files_test.go

@@ -25,3 +25,10 @@ func TestWriteBadlyFormattedBase64DataURI(t *testing.T) {
 	_, err := WriteDataURLToFile("data:None;base64,SGVsbG8gVGhlcmU=", path)
 	require.NoError(t, err)
 }
+
+func TestWriteNotRecognisedBase64DataURL(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test-file")
+	_, err := WriteDataURLToFile("data:None;model/gltf-binary,SGVsbG8gVGhlcmU=", path)
+	require.NoError(t, err)
+}

test-integration/test_integration/test_build.py

@@ -595,3 +595,14 @@ def test_torch_271_cuda_128_base_image(docker_image, cog_binary):
         capture_output=True,
     )
     assert build_process.returncode == 0
+
+
+def test_python_313_base_images(docker_image, cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/python-313"
+
+    build_process = subprocess.run(
+        [cog_binary, "build", "-t", docker_image, "--use-cog-base-image"],
+        cwd=project_dir,
+        capture_output=True,
+    )
+    assert build_process.returncode == 0