refactor: use rest/grpc for the wait strategy

camunda-community-hub · Dec 29, 2024 · efe5abd · efe5abd
1 parent 72579c8
commit efe5abd
Showing 1 changed file with 13 additions and 6 deletions.
diff --git a/core/src/main/java/io/zeebe/containers/ZeebeTopologyWaitStrategy.java b/core/src/main/java/io/zeebe/containers/ZeebeTopologyWaitStrategy.java
@@ -22,6 +22,7 @@
 import io.camunda.zeebe.client.api.response.BrokerInfo;
 import io.camunda.zeebe.client.api.response.PartitionInfo;
 import io.camunda.zeebe.client.api.response.Topology;
+
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -35,7 +36,6 @@
 import org.slf4j.LoggerFactory;
 import org.testcontainers.containers.ContainerLaunchException;
 import org.testcontainers.containers.wait.strategy.AbstractWaitStrategy;
-import org.testcontainers.containers.wait.strategy.WaitStrategyTarget;
 
 /**
  * A {@link org.testcontainers.containers.wait.strategy.WaitStrategy} implementation which waits for
@@ -200,9 +200,17 @@ public ZeebeTopologyWaitStrategy forBuilder(
 
   @Override
   protected void waitUntilReady() {
+    if (!(waitStrategyTarget instanceof ZeebeGatewayNode<?>)) {
+      throw new UnsupportedOperationException(
+          String.format(
+              "Expected to await a ZeebeGatewayNode, but given target is a %s; cannot proceed",
+              waitStrategyTarget.getClass()));
+    }
+
+    final ZeebeGatewayNode<?> node = (ZeebeGatewayNode<?>) waitStrategyTarget;
     final TopologyHolder latestTopology = new TopologyHolder();
 
-    try (final ZeebeClient client = newZeebeClient(waitStrategyTarget)) {
+    try (final ZeebeClient client = newZeebeClient(node)) {
       final String containerName = waitStrategyTarget.getContainerInfo().getName();
       LOGGER.info(
           "{}: Waiting for {} for topology to have at least {} brokers, {} partitions with "
@@ -232,12 +240,11 @@ protected void waitUntilReady() {
     }
   }
 
-  private ZeebeClient newZeebeClient(final WaitStrategyTarget waitStrategyTarget) {
-    final String gatewayHost = waitStrategyTarget.getHost();
-    final int exposedGatewayPort = waitStrategyTarget.getMappedPort(gatewayPort);
+  private ZeebeClient newZeebeClient(final ZeebeGatewayNode<?> gateway) {
     return clientBuilderProvider
         .get()
-        .gatewayAddress(gatewayHost + ":" + exposedGatewayPort)
+        .grpcAddress(gateway.getGrpcAddress())
+        .restAddress(gateway.getRestAddress())
         .build();
   }