Merge pull request #143 from aivinog1/139-set-image-name-in-builder

feat: add the ability to set specific images via ZeebeClusterBuilder

src/main/java/io/zeebe/containers/cluster/ZeebeClusterBuilder.java

@@ -15,6 +15,8 @@
  */
 package io.zeebe.containers.cluster;
 
+import static io.zeebe.containers.ZeebeDefaults.getInstance;
+
 import io.zeebe.containers.ZeebeBrokerContainer;
 import io.zeebe.containers.ZeebeBrokerNode;
 import io.zeebe.containers.ZeebeContainer;
@@ -31,6 +33,7 @@
 import org.apiguardian.api.API.Status;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.Network;
+import org.testcontainers.utility.DockerImageName;
 
 /**
  * Convenience class to help build Zeebe clusters.
@@ -104,6 +107,8 @@ public class ZeebeClusterBuilder {
   private int partitionsCount = 1;
   private int replicationFactor = 1;
   private boolean useEmbeddedGateway = true;
+  private DockerImageName gatewayImageName = getInstance().getDefaultDockerImage();
+  private DockerImageName brokerImageName = getInstance().getDefaultDockerImage();
 
   private final Map<String, ZeebeGatewayNode<? extends GenericContainer<?>>> gateways =
       new HashMap<>();
@@ -242,6 +247,41 @@ public ZeebeClusterBuilder withName(final String name) {
     return this;
   }
 
+  /**
+   * Sets the docker image for separate gateways. This could be used to create a test against
+   * specific Zeebe version.
+   *
+   * @param gatewayImageName the docker image name of the Zeebe
+   * @return this builder instance for chaining
+   */
+  public ZeebeClusterBuilder withGatewayImage(final DockerImageName gatewayImageName) {
+    this.gatewayImageName = Objects.requireNonNull(gatewayImageName);
+    return this;
+  }
+
+  /**
+   * Sets the docker image for brokers with embedded gateways. This could be used to create a test
+   * against specific Zeebe version.
+   *
+   * @param brokerImageName the docker image name of the Zeebe
+   * @return this builder instance for chaining
+   */
+  public ZeebeClusterBuilder withBrokerImage(final DockerImageName brokerImageName) {
+    this.brokerImageName = Objects.requireNonNull(brokerImageName);
+    return this;
+  }
+
+  /**
+   * Sets the docker image for brokers with embedded gateways and separate gateways. This could be
+   * used to create a test against specific Zeebe version.
+   *
+   * @param imageName the docker image name of the Zeebe
+   * @return this builder instance for chaining
+   */
+  public ZeebeClusterBuilder withImage(final DockerImageName imageName) {
+    return withGatewayImage(imageName).withBrokerImage(imageName);
+  }
+
   /**
    * Builds a new Zeebe cluster. Will create {@link #brokersCount} brokers (accessible later via
    * {@link ZeebeCluster#getBrokers()}) and {@link #gatewaysCount} standalone gateways (accessible
@@ -305,13 +345,13 @@ private void createBrokers() {
       final ZeebeBrokerNode<?> broker;
 
       if (useEmbeddedGateway) {
-        final ZeebeContainer container = new ZeebeContainer();
+        final ZeebeContainer container = new ZeebeContainer(brokerImageName);
         configureGateway(container);
 
         broker = container;
         gateways.put(String.valueOf(i), container);
       } else {
-        broker = new ZeebeBrokerContainer();
+        broker = new ZeebeBrokerContainer(brokerImageName);
       }
 
       configureBroker(broker, i);
@@ -341,7 +381,7 @@ private void createStandaloneGateways() {
   }
 
   private ZeebeGatewayContainer createStandaloneGateway(final String memberId) {
-    final ZeebeGatewayContainer gateway = new ZeebeGatewayContainer();
+    final ZeebeGatewayContainer gateway = new ZeebeGatewayContainer(gatewayImageName);
 
     gateway
         .withNetwork(network)

src/test/java/io/zeebe/containers/cluster/ZeebeClusterBuilderTest.java

@@ -25,10 +25,12 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import org.assertj.core.api.Condition;
 import org.junit.jupiter.api.Test;
 import org.testcontainers.containers.Container;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.Network;
+import org.testcontainers.utility.DockerImageName;
 
 final class ZeebeClusterBuilderTest {
   @Test
@@ -506,4 +508,81 @@ void shouldNotConfigureEmbeddedGateway() {
         .as("the broker is not configured to enable the embedded gateway")
         .doesNotContainEntry("ZEEBE_BROKER_GATEWAY_ENABLE", "true");
   }
+
+  @Test
+  void shouldSetImageNameForGateways() {
+    // given
+    final ZeebeClusterBuilder builder = new ZeebeClusterBuilder();
+    final String zeebeDockerImage = "camunda/zeebe:latest";
+
+    // when
+    final DockerImageName gatewayImageName = DockerImageName.parse(zeebeDockerImage);
+    builder.withGatewayImage(gatewayImageName).withGatewaysCount(1).withEmbeddedGateway(false);
+    final ZeebeCluster cluster = builder.build();
+
+    // then
+    assertThat(cluster.getGateways().entrySet())
+        .as("the only gateway created has the right docker image")
+        .singleElement()
+        .satisfies(
+            gatewayEntry -> verifyZeebeHasImageName(gatewayEntry.getValue(), zeebeDockerImage));
+  }
+
+  @Test
+  void shouldSetImageNameForBrokers() {
+    // given
+    final ZeebeClusterBuilder builder = new ZeebeClusterBuilder();
+    final String zeebeDockerImage = "camunda/zeebe:latest";
+
+    // when
+    final DockerImageName gatewayImageName = DockerImageName.parse(zeebeDockerImage);
+    builder.withBrokerImage(gatewayImageName).withBrokersCount(1);
+    final ZeebeCluster cluster = builder.build();
+
+    // then
+    assertThat(cluster.getBrokers().entrySet())
+        .as("the only broker created has the right docker image")
+        .singleElement()
+        .satisfies(
+            brokerEntry -> verifyZeebeHasImageName(brokerEntry.getValue(), zeebeDockerImage));
+  }
+
+  @Test
+  void shouldSetImageNameForGatewaysAndBrokers() {
+    // given
+    final ZeebeClusterBuilder builder = new ZeebeClusterBuilder();
+    final String zeebeDockerImage = "camunda/zeebe:latest";
+
+    // when
+    final DockerImageName gatewayImageName = DockerImageName.parse(zeebeDockerImage);
+    builder
+        .withImage(gatewayImageName)
+        .withBrokersCount(1)
+        .withGatewaysCount(1)
+        .withEmbeddedGateway(false);
+    final ZeebeCluster cluster = builder.build();
+
+    // then
+    assertThat(cluster.getBrokers().entrySet())
+        .as("the only broker created has the right docker image")
+        .singleElement()
+        .satisfies(
+            brokerEntry -> verifyZeebeHasImageName(brokerEntry.getValue(), zeebeDockerImage));
+    assertThat(cluster.getGateways().entrySet())
+        .as("the only standalone gateway created has the right docker image")
+        .singleElement()
+        .satisfies(
+            gatewayEntry -> verifyZeebeHasImageName(gatewayEntry.getValue(), zeebeDockerImage));
+  }
+
+  private Condition<ZeebeNode<? extends GenericContainer<?>>> zeebeImageHasImageName(
+      final String imageName) {
+    return new Condition<>(
+        node -> node.getDockerImageName().equals(imageName), "Image Name Condition");
+  }
+
+  private void verifyZeebeHasImageName(
+      final ZeebeNode<? extends GenericContainer<?>> zeebe, final String imageName) {
+    assertThat(zeebe.getDockerImageName()).isEqualTo(imageName);
+  }
 }