chore: improvements for test repeatability (#766)

Signed-off-by: Alex Kehayov <[email protected]>
hiero-ledger · Mar 11, 2025 · 14f1d62 · 14f1d62
1 parent c42d534
commit 14f1d62
Show file tree

Hide file tree

Showing 5 changed files with 241 additions and 84 deletions.
diff --git a/server/src/main/java/com/hedera/block/server/pbj/PbjBlockStreamServiceProxy.java b/server/src/main/java/com/hedera/block/server/pbj/PbjBlockStreamServiceProxy.java
@@ -130,7 +130,37 @@ public Pipeline<? super Bytes> open(
             final @NonNull RequestOptions options,
             final @NonNull Pipeline<? super Bytes> replies) {
 
-        final var m = (BlockStreamMethod) method;
+        return openPipeline(
+                (BlockStreamMethod) method,
+                options,
+                replies,
+                openRangeHistoricStreamingExecutorService,
+                closedRangeHistoricStreamingExecutorService);
+    }
+
+    @NonNull
+    public Pipeline<? super Bytes> open(
+            final @NonNull Method method,
+            final @NonNull RequestOptions options,
+            final @NonNull Pipeline<? super Bytes> replies,
+            final @NonNull ExecutorService openRangeHistoricStreamingExecutorService,
+            final @NonNull ExecutorService closedRangeHistoricStreamingExecutorService) {
+
+        return openPipeline(
+                (BlockStreamMethod) method,
+                options,
+                replies,
+                openRangeHistoricStreamingExecutorService,
+                closedRangeHistoricStreamingExecutorService);
+    }
+
+    private Pipeline<? super Bytes> openPipeline(
+            @NonNull BlockStreamMethod method,
+            @NonNull RequestOptions options,
+            @NonNull Pipeline<? super Bytes> replies,
+            @NonNull ExecutorService openRangeHistoricStreamingExecutorService,
+            @NonNull ExecutorService closedRangeHistoricStreamingExecutorService) {
+        final var m = method;
         try {
             return switch (m) {
                 case publishBlockStream -> {
@@ -145,7 +175,11 @@ public Pipeline<? super Bytes> open(
                 case subscribeBlockStream -> Pipelines
                         .<SubscribeStreamRequest, SubscribeStreamResponseUnparsed>serverStreaming()
                         .mapRequest(bytes -> parseSubscribeStreamRequest(bytes, options))
-                        .method(this::subscribeBlockStream)
+                        .method((subscribeStreamRequest, helidonConsumerObserver) -> subscribeBlockStream(
+                                subscribeStreamRequest,
+                                helidonConsumerObserver,
+                                openRangeHistoricStreamingExecutorService,
+                                closedRangeHistoricStreamingExecutorService))
                         .mapResponse(reply -> createSubscribeStreamResponse(reply, options))
                         .respondTo(replies)
                         .build();
@@ -201,7 +235,9 @@ Pipeline<List<BlockItemUnparsed>> publishBlockStream(
      */
     void subscribeBlockStream(
             @NonNull final SubscribeStreamRequest subscribeStreamRequest,
-            @NonNull final Pipeline<? super SubscribeStreamResponseUnparsed> helidonConsumerObserver) {
+            @NonNull final Pipeline<? super SubscribeStreamResponseUnparsed> helidonConsumerObserver,
+            @NonNull final ExecutorService openRangeHistoricStreamingExecutorService,
+            @NonNull final ExecutorService closedRangeHistoricStreamingExecutorService) {
 
         LOGGER.log(DEBUG, "Executing Server Streaming subscribeBlockStream gRPC method");
 

diff --git a/server/src/test/java/com/hedera/block/server/mediator/LiveStreamMediatorImplTest.java b/server/src/test/java/com/hedera/block/server/mediator/LiveStreamMediatorImplTest.java
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.hedera.block.server.mediator;
 
+import static com.hedera.block.server.metrics.BlockNodeMetricTypes.Counter.BlocksPersisted;
 import static com.hedera.block.server.metrics.BlockNodeMetricTypes.Counter.LiveBlockItems;
 import static com.hedera.block.server.metrics.BlockNodeMetricTypes.Gauge.Consumers;
 import static com.hedera.block.server.util.PersistTestUtils.PERSISTENCE_STORAGE_ARCHIVE_ROOT_PATH_KEY;
@@ -32,6 +33,7 @@
 import com.hedera.block.server.service.ServiceConfig;
 import com.hedera.block.server.service.ServiceStatus;
 import com.hedera.block.server.service.ServiceStatusImpl;
+import com.hedera.block.server.util.BlockingExecutorService;
 import com.hedera.block.server.util.PersistTestUtils;
 import com.hedera.block.server.util.TestConfigUtil;
 import com.hedera.hapi.block.BlockItemSetUnparsed;
@@ -174,7 +176,9 @@ void testUnsubscribeEach() throws InterruptedException {
     }
 
     @Test
-    void testMediatorPersistenceWithoutSubscribers() throws IOException {
+    void testMediatorPersistenceWithoutSubscribers() throws IOException, InterruptedException {
+        // 1 block is expected to be processed, so the expected tasks param is set to 1
+        final BlockingExecutorService executor = new BlockingExecutorService(1, 1);
         final ServiceStatus serviceStatus = new ServiceStatusImpl(serviceConfig);
         final LiveStreamMediator streamMediator = LiveStreamMediatorBuilder.newBuilder(
                         metricsService, mediatorConfig, serviceStatus)
@@ -192,7 +196,7 @@ void testMediatorPersistenceWithoutSubscribers() throws IOException {
                 serviceStatus,
                 ackHandlerMock,
                 writerFactory,
-                executorMock,
+                executor,
                 archiverMock,
                 pathResolverMock,
                 persistenceStorageConfig);
@@ -201,13 +205,12 @@ void testMediatorPersistenceWithoutSubscribers() throws IOException {
         // Acting as a producer, notify the mediator of a new block
         streamMediator.publish(blockItemUnparsed);
 
+        // Wait all the tasks to complete before the assertions to avoid flakiness
+        executor.waitTasksToComplete();
+
         // Verify the counter was incremented
         assertEquals(10, metricsService.get(LiveBlockItems).get());
-
-        // @todo(642) we need to employ the same technique here to inject a writer that will ensure
-        // the tasks are complete before we can verify the metrics for blocks persisted
-        // the test will pass without flaking if we do that
-        // assertEquals(1, blockNodeContext.metricsService().get(BlocksPersisted).get());
+        assertEquals(1, metricsService.get(BlocksPersisted).get());
     }
 
     @Test