From 3cd79ca9bf087d1f044e239fa503232c87d99831 Mon Sep 17 00:00:00 2001
From: "Andy Su (Apps)" <asu80@bloomberg.net>
Date: Tue, 27 Feb 2024 16:48:41 -0500
Subject: [PATCH] allow customization of test containers image name

---
 .../gateway/CustomImageNameSubstitutor.java   | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)
 create mode 100644 gateway-ha/src/test/java/io/trino/gateway/CustomImageNameSubstitutor.java

diff --git a/gateway-ha/src/test/java/io/trino/gateway/CustomImageNameSubstitutor.java b/gateway-ha/src/test/java/io/trino/gateway/CustomImageNameSubstitutor.java
new file mode 100644
index 000000000..5f6be9074
--- /dev/null
+++ b/gateway-ha/src/test/java/io/trino/gateway/CustomImageNameSubstitutor.java
@@ -0,0 +1,25 @@
+package io.trino.gateway;
+
+import org.testcontainers.utility.DockerImageName;
+import org.testcontainers.utility.ImageNameSubstitutor;
+
+public class CustomImageNameSubstitutor
+        extends ImageNameSubstitutor
+{
+    private final String TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE = "TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE";
+    @Override
+    public DockerImageName apply(DockerImageName dockerImageName)
+    {
+        if (dockerImageName.asCanonicalNameString().equals("trinodb/trino:latest") && System.getenv(TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE) != null) {
+            String image_substitute = System.getenv(TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE);
+            return DockerImageName.parse(image_substitute).asCompatibleSubstituteFor("trinodb/trino");
+        }
+        return dockerImageName;
+    }
+
+    @Override
+    protected String getDescription()
+    {
+        return "custom image name substitutor";
+    }
+}