refactor(engine): implement new API

camunda-community-hub · Dec 20, 2023 · 26cb62a · 26cb62a
1 parent 81afe98
commit 26cb62a
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 7 deletions.
diff --git a/engine/src/main/java/io/zeebe/containers/engine/ZeebeClusterEngine.java b/engine/src/main/java/io/zeebe/containers/engine/ZeebeClusterEngine.java
@@ -15,7 +15,10 @@
  */
 package io.zeebe.containers.engine;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import io.camunda.zeebe.client.ZeebeClient;
+import io.camunda.zeebe.client.ZeebeClientBuilder;
+import io.camunda.zeebe.client.impl.ZeebeObjectMapper;
 import io.camunda.zeebe.process.test.api.RecordStreamSource;
 import io.zeebe.containers.ZeebeNode;
 import io.zeebe.containers.clock.ZeebeClock;
@@ -25,6 +28,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.TimeoutException;
+import java.util.function.UnaryOperator;
 import org.agrona.CloseHelper;
 import org.apiguardian.api.API;
 import org.apiguardian.api.API.Status;
@@ -65,10 +69,12 @@ public RecordStreamSource getRecordStreamSource() {
 
   @Override
   public ZeebeClient createClient() {
-    final ZeebeClient client = cluster.newClientBuilder().build();
-    clients.add(client);
+    return createClient(UnaryOperator.identity());
+  }
 
-    return client;
+  @Override
+  public ZeebeClient createClient(final ObjectMapper customObjectMapper) {
+    return createClient(b -> b.withJsonMapper(new ZeebeObjectMapper(customObjectMapper)));
   }
 
   @Override
@@ -106,4 +112,14 @@ public void stop() {
 
     CloseHelper.closeAll(cluster, recordStream);
   }
+
+  private ZeebeClient createClient(final UnaryOperator<ZeebeClientBuilder> configurator) {
+    final ZeebeClientBuilder builder =
+        configurator.apply(
+            ZeebeClient.newClientBuilder().usePlaintext().gatewayAddress(getGatewayAddress()));
+    final ZeebeClient client = builder.build();
+    clients.add(client);
+
+    return client;
+  }
 }
diff --git a/engine/src/main/java/io/zeebe/containers/engine/ZeebeContainerEngine.java b/engine/src/main/java/io/zeebe/containers/engine/ZeebeContainerEngine.java
@@ -15,7 +15,10 @@
  */
 package io.zeebe.containers.engine;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import io.camunda.zeebe.client.ZeebeClient;
+import io.camunda.zeebe.client.ZeebeClientBuilder;
+import io.camunda.zeebe.client.impl.ZeebeObjectMapper;
 import io.camunda.zeebe.process.test.api.RecordStreamSource;
 import io.zeebe.containers.ZeebeBrokerNode;
 import io.zeebe.containers.ZeebeContainer;
@@ -26,6 +29,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.TimeoutException;
+import java.util.function.UnaryOperator;
 import org.agrona.CloseHelper;
 import org.apiguardian.api.API;
 import org.apiguardian.api.API.Status;
@@ -63,11 +67,12 @@ public RecordStreamSource getRecordStreamSource() {
 
   @Override
   public ZeebeClient createClient() {
-    final ZeebeClient client =
-        ZeebeClient.newClientBuilder().usePlaintext().gatewayAddress(getGatewayAddress()).build();
-    clients.add(client);
+    return createClient(UnaryOperator.identity());
+  }
 
-    return client;
+  @Override
+  public ZeebeClient createClient(final ObjectMapper objectMapper) {
+    return createClient(b -> b.withJsonMapper(new ZeebeObjectMapper(objectMapper)));
   }
 
   @Override
@@ -105,4 +110,14 @@ public void stop() {
 
     CloseHelper.closeAll(container, recordStream);
   }
+
+  private ZeebeClient createClient(final UnaryOperator<ZeebeClientBuilder> configurator) {
+    final ZeebeClientBuilder builder =
+        configurator.apply(
+            ZeebeClient.newClientBuilder().usePlaintext().gatewayAddress(getGatewayAddress()));
+    final ZeebeClient client = builder.build();
+    clients.add(client);
+
+    return client;
+  }
 }