Allow customization of test containers image name

gateway-ha/src/test/java/io/trino/gateway/CustomTrinoImageNameSubstitutor.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.gateway;
+
+import org.testcontainers.utility.DockerImageName;
+import org.testcontainers.utility.ImageNameSubstitutor;
+
+public class CustomTrinoImageNameSubstitutor
+        extends ImageNameSubstitutor
+{
+    private static final String TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE = "TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE";
+
+    @Override
+    public DockerImageName apply(DockerImageName dockerImageName)
+    {
+        String trinoImage = "trinodb/trino";
+        if (dockerImageName.getUnversionedPart().equals(trinoImage) && System.getenv(TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE) != null) {
+            String imageSubstitute = System.getenv(TESTCONTAINERS_TRINO_IMAGE_SUBSTITUTE);
+            return DockerImageName.parse(imageSubstitute).asCompatibleSubstituteFor(trinoImage);
+        }
+        return dockerImageName;
+    }
+
+    @Override
+    protected String getDescription()
+    {
+        return "custom image name substitutor";
+    }
+}