From 8854c9d5a2d45e9d3e1cd040bb3f774670bc765b Mon Sep 17 00:00:00 2001
From: Nicolas Pepin-Perreault <nicolas.pepin-perreault@camunda.com>
Date: Sun, 29 Dec 2024 22:17:18 +0100
Subject: [PATCH] refactor: use rest/grpc for the wait strategy

---
 .../containers/ZeebeTopologyWaitStrategy.java    | 16 +++++++++++++---
 .../ZeebeTopologyWaitStrategyTest.java           |  9 ++++++---
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/core/src/main/java/io/zeebe/containers/ZeebeTopologyWaitStrategy.java b/core/src/main/java/io/zeebe/containers/ZeebeTopologyWaitStrategy.java
index ff80e4d9..64708ec6 100644
--- 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.net.URI;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -233,11 +234,20 @@ protected void waitUntilReady() {
   }
 
   private ZeebeClient newZeebeClient(final WaitStrategyTarget waitStrategyTarget) {
-    final String gatewayHost = waitStrategyTarget.getHost();
-    final int exposedGatewayPort = waitStrategyTarget.getMappedPort(gatewayPort);
+    //noinspection HttpUrlsUsage
+    final URI gatewayAddress =
+        URI.create(
+            "http://"
+                + waitStrategyTarget.getHost()
+                + ":"
+                + waitStrategyTarget.getMappedPort(gatewayPort));
+
+    // use the same URI for the REST and gRPC address, and rely on the user to correctly set the
+    // port they wish to use
     return clientBuilderProvider
         .get()
-        .gatewayAddress(gatewayHost + ":" + exposedGatewayPort)
+        .restAddress(gatewayAddress)
+        .grpcAddress(gatewayAddress)
         .build();
   }
 
diff --git a/core/src/test/java/io/zeebe/containers/ZeebeTopologyWaitStrategyTest.java b/core/src/test/java/io/zeebe/containers/ZeebeTopologyWaitStrategyTest.java
index 7b7607fe..5f559270 100644
--- a/core/src/test/java/io/zeebe/containers/ZeebeTopologyWaitStrategyTest.java
+++ b/core/src/test/java/io/zeebe/containers/ZeebeTopologyWaitStrategyTest.java
@@ -26,6 +26,7 @@
 import io.camunda.zeebe.gateway.protocol.GatewayOuterClass.Partition;
 import io.camunda.zeebe.gateway.protocol.GatewayOuterClass.Partition.PartitionBrokerRole;
 import io.camunda.zeebe.gateway.protocol.GatewayOuterClass.TopologyResponse;
+import java.net.URI;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
@@ -70,7 +71,8 @@ final class ZeebeTopologyWaitStrategyTest {
 
   @BeforeEach
   void setup() {
-    Mockito.when(builder.gatewayAddress(Mockito.anyString())).thenReturn(builder);
+    Mockito.when(builder.grpcAddress(Mockito.any())).thenReturn(builder);
+    Mockito.when(builder.restAddress(Mockito.any())).thenReturn(builder);
     Mockito.when(builder.build()).thenReturn(client);
     Mockito.when(client.newTopologyRequest()).thenReturn(topologyRequest);
   }
@@ -132,8 +134,9 @@ void shouldUseCorrectGatewayAddress() {
     strategy.waitUntilReady(target);
 
     // then
-    Mockito.verify(builder, Mockito.timeout(5000).atLeastOnce())
-        .gatewayAddress(target.getHost() + ":" + target.mappedPort);
+    final URI expectedAddress = URI.create("http://" + target.getHost() + ":" + target.mappedPort);
+    Mockito.verify(builder, Mockito.timeout(5000).atLeastOnce()).grpcAddress(expectedAddress);
+    Mockito.verify(builder, Mockito.timeout(5000).atLeastOnce()).restAddress(expectedAddress);
   }
 
   @ParameterizedTest(name = "should timeout on incomplete topology when {0}")