Axon test util for waiting for some events

AxonIntegrationTest.java

class AxonIntegrationTest extends AbstractIntegrationTest {
 
  @Test
  void create_document_in_google_docs__propagate_to_ms_word_online(@Autowired EventBus eventBus){
    AxonTestUtil.executeAndWaitForEvents(
            () -> documentService.createGoogleDocument(...),
            eventBus, 2, DocumentSavedEvent.class
    );
    //instead of this:
    //documentService.createGoogleDocument(...);
    //TimeUnit.SECONDS.sleep(1);
    verify(googleDocsClient).create(any());
    verify(msWordOnlineClient).create(any());
  }
}

AxonTestUtil.java

package de.applink.processing.engine;

import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import lombok.val;
import org.axonframework.eventhandling.EventBus;
import org.axonframework.eventhandling.EventMessage;
import org.axonframework.messaging.MessageDispatchInterceptor;

import java.util.concurrent.atomic.AtomicInteger;

import static org.assertj.core.api.Assertions.assertThat;
import static org.awaitility.Awaitility.await;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;

@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class AxonTestUtil {

    public static void executeAndWaitForEvents(Runnable action, EventBus eventBus, int amount, Class<?> eventClass) {
        val foundEvents = executeAndWaitForEventsAtLeast(action, eventBus, amount, eventClass);
        assertThat(foundEvents).as(amount + " " + eventClass.getSimpleName() + " expected").isEqualTo(amount);
    }

    /**
     * @return amount of events found. Greater or equal to the passed amount.
     */
    //this could also accept Map<Integer, Class<?>> to count multiple events
    public static int executeAndWaitForEventsAtLeast(Runnable action, EventBus eventBus, int amount, Class<?> eventClass) {
        val eventCounter = new AtomicInteger();
        val registration = eventBus.registerDispatchInterceptor(eventCounterInterceptor(eventCounter, eventClass));

        action.run();
        await(amount + " " + eventClass.getSimpleName() + " expected").untilAtomic(eventCounter, greaterThanOrEqualTo(amount));
        registration.cancel();
        return eventCounter.intValue();
    }

    private static MessageDispatchInterceptor<EventMessage<?>> eventCounterInterceptor(AtomicInteger i, Class<?> eventClass) {
        return messages -> (index, event) -> {
            if (event.getPayloadType().equals(eventClass)) {
                i.incrementAndGet();
            }
            return event;
        };
    }
}