ok test

transferwise · Aug 6, 2024 · 52a348d · 52a348d
1 parent 7eab357
commit 52a348d
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tw-tkms-starter/src/test/java/com/transferwise/kafka/tkms/test/TestMessageDecorator.java b/tw-tkms-starter/src/test/java/com/transferwise/kafka/tkms/test/TestMessageDecorator.java
@@ -13,16 +13,18 @@ public class TestMessageDecorator implements ITkmsMessageDecorator {
 
   @Override
   public List<Header> getAdditionalHeaders(TkmsMessage message) {
-    var h1 = new Header().setKey("tool").setValue("jambi".getBytes(StandardCharsets.UTF_8));
     if (message.getValue() != null && new String(message.getValue(), StandardCharsets.UTF_8).startsWith("Here from")) {
-      return List.of(h1);
+      return List.of(new Header().setKey("tool").setValue("jambi".getBytes(StandardCharsets.UTF_8)));
     }
     return List.of();
   }
 
   @Override
   public TkmsShardPartition getOverridedPartition(TkmsMessage message) {
-    return new TkmsShardPartition(0, 0);
+    if (message.getValue() != null && new String(message.getValue(), StandardCharsets.UTF_8).startsWith("Here from")) {
+      return new TkmsShardPartition(0, 0);
+    }
+    return new TkmsShardPartition(message.getShard(), message.getPartition());
   }
 
 }