[#2154] Fix missing propagation of ReplayToken to UnitOfWork for Pool…

…edStreamingEventProcessor
AxonFramework · Mar 28, 2022 · 389b6a6 · 389b6a6
1 parent ab61170
commit 389b6a6
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/messaging/src/main/java/org/axonframework/eventhandling/pooled/WorkPackage.java b/messaging/src/main/java/org/axonframework/eventhandling/pooled/WorkPackage.java
@@ -228,7 +228,10 @@ private void processEvents() throws Exception {
             ProcessingEntry entry = processingQueue.poll();
             lastConsumedToken = WrappedToken.advance(lastConsumedToken, entry.eventMessage().trackingToken());
             if (entry.canHandle()) {
-                eventBatch.add(entry.eventMessage());
+                eventBatch.add(
+                        entry.eventMessage()
+                             .withTrackingToken(lastConsumedToken)
+                );
             }
         }
 

diff --git a/messaging/src/test/java/org/axonframework/eventhandling/pooled/WorkPackageTest.java b/messaging/src/test/java/org/axonframework/eventhandling/pooled/WorkPackageTest.java
@@ -171,7 +171,7 @@ void testScheduleEventFailsOnBatchProcessor() throws ExecutionException, Interru
         TrackedEventMessage<String> testEvent =
                 new GenericTrackedEventMessage<>(testToken, GenericEventMessage.asEventMessage("some-event"));
         batchProcessorPredicate = event -> {
-            if (event.contains(testEvent)) {
+            if (event.stream().anyMatch(e -> ((TrackedEventMessage<?>) e).trackingToken().equals(testToken))) {
                 throw new IllegalStateException("Some exception");
             }
             return true;
@@ -207,7 +207,7 @@ void testScheduleEventRunsSuccessfully() {
 
         List<EventMessage<?>> processedEvents = batchProcessor.getProcessedEvents();
         assertWithin(500, TimeUnit.MILLISECONDS, () -> assertEquals(1, processedEvents.size()));
-        assertEquals(expectedEvent, processedEvents.get(0));
+        assertEquals(expectedEvent.trackingToken(), ((TrackedEventMessage<?>) processedEvents.get(0)).trackingToken());
 
         ArgumentCaptor<TrackingToken> tokenCaptor = ArgumentCaptor.forClass(TrackingToken.class);
         verify(tokenStore).storeToken(tokenCaptor.capture(), eq(PROCESSOR_NAME), eq(segment.getSegmentId()));
@@ -236,7 +236,7 @@ void testScheduleEventIsScheduledIfTheLastDeliveredTokenEqualsTheEventsToken() {
 
         List<EventMessage<?>> processedEvents = batchProcessor.getProcessedEvents();
         assertWithin(500, TimeUnit.MILLISECONDS, () -> assertEquals(1, processedEvents.size()));
-        assertEquals(expectedEvent, processedEvents.get(0));
+        assertEquals(expectedEvent.trackingToken(), ((TrackedEventMessage<?>) processedEvents.get(0)).trackingToken());
 
         ArgumentCaptor<TrackingToken> tokenCaptor = ArgumentCaptor.forClass(TrackingToken.class);
         verify(tokenStore).storeToken(tokenCaptor.capture(), eq(PROCESSOR_NAME), eq(segment.getSegmentId()));
@@ -264,7 +264,7 @@ void testScheduleEventExtendsTokenClaimAfterClaimThresholdExtension() {
         // Should have handled one event, so a subsequent run of WorkPackage#processEvents will extend the claim.
         List<EventMessage<?>> processedEvents = batchProcessor.getProcessedEvents();
         assertWithin(500, TimeUnit.MILLISECONDS, () -> assertEquals(1, processedEvents.size()));
-        assertEquals(expectedEvent, processedEvents.get(0));
+        assertEquals(expectedEvent.trackingToken(), ((TrackedEventMessage<?>) processedEvents.get(0)).trackingToken());
         // We  need to verify the TokenStore#storeToken operation, otherwise the extendClaim verify will not succeed.
         ArgumentCaptor<TrackingToken> tokenCaptor = ArgumentCaptor.forClass(TrackingToken.class);
         verify(tokenStore).storeToken(tokenCaptor.capture(), eq(PROCESSOR_NAME), eq(segment.getSegmentId()));
@@ -416,4 +416,4 @@ public List<EventMessage<?>> getProcessedEvents() {
             return processedEvents;
         }
     }
-}
+}