fix: injected map to make it more testable

Signed-off-by: Matt Peterson <[email protected]>
hashgraph · Jul 24, 2024 · 6988dde · 6988dde
1 parent 7916f42
commit 6988dde
Show file tree

Hide file tree

Showing 4 changed files with 310 additions and 28 deletions.
diff --git a/server/src/main/java/com/hedera/block/server/consumer/ConsumerBlockItemObserver.java b/server/src/main/java/com/hedera/block/server/consumer/ConsumerBlockItemObserver.java
@@ -16,15 +16,15 @@
 
 package com.hedera.block.server.consumer;
 
-import static com.hedera.block.protos.BlockStreamService.BlockItem;
-import static com.hedera.block.protos.BlockStreamService.SubscribeStreamResponse;
-
 import com.hedera.block.server.data.ObjectEvent;
 import com.hedera.block.server.mediator.StreamMediator;
 import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
+
 import java.time.InstantSource;
-import java.util.concurrent.CountDownLatch;
+
+import static com.hedera.block.protos.BlockStreamService.BlockItem;
+import static com.hedera.block.protos.BlockStreamService.SubscribeStreamResponse;
 
 /**
  * The LiveStreamObserverImpl class implements the LiveStreamObserver interface to pass blocks to
@@ -41,8 +41,6 @@ public class ConsumerBlockItemObserver implements BlockItemEventHandler<ObjectEv
     private final InstantSource producerLivenessClock;
     private long producerLivenessMillis;
 
-    private final CountDownLatch shutdownLatch = new CountDownLatch(1);
-
     private final StreamMediator<ObjectEvent<BlockItem>, BlockItem> streamMediator;
 
     private boolean streamStarted;
@@ -118,6 +116,5 @@ public void onEvent(final ObjectEvent<BlockItem> event, final long l, final bool
 
     @Override
     public void awaitShutdown() throws InterruptedException {
-        shutdownLatch.await();
     }
 }
diff --git a/server/src/main/java/com/hedera/block/server/mediator/LiveStreamMediatorImpl.java b/server/src/main/java/com/hedera/block/server/mediator/LiveStreamMediatorImpl.java
@@ -50,7 +50,7 @@ public class LiveStreamMediatorImpl implements StreamMediator<ObjectEvent<BlockI
     private final Map<
                     BlockItemEventHandler<ObjectEvent<BlockItem>>,
                     BatchEventProcessor<ObjectEvent<BlockItem>>>
-            subscribers = new HashMap<>();
+            subscribers;
 
     private final BlockPersistenceHandler<BlockItem, Block> blockPersistenceHandler;
     private final Consumer<StreamMediator<ObjectEvent<BlockItem>, BlockItem>> shutdownCallback;
@@ -61,8 +61,13 @@ public class LiveStreamMediatorImpl implements StreamMediator<ObjectEvent<BlockI
      * @param blockPersistenceHandler the block persistence handler
      */
     public LiveStreamMediatorImpl(
+            final Map<
+                            BlockItemEventHandler<ObjectEvent<BlockItem>>,
+                            BatchEventProcessor<ObjectEvent<BlockItem>>>
+                    subscribers,
             final BlockPersistenceHandler<BlockItem, Block> blockPersistenceHandler,
             final Consumer<StreamMediator<ObjectEvent<BlockItem>, BlockItem>> shutdownCallback) {
+        this.subscribers = subscribers;
         this.blockPersistenceHandler = blockPersistenceHandler;
 
         // Initialize and start the disruptor
@@ -73,6 +78,12 @@ public LiveStreamMediatorImpl(
         this.shutdownCallback = shutdownCallback;
     }
 
+    public LiveStreamMediatorImpl(
+            final BlockPersistenceHandler<BlockItem, Block> blockPersistenceHandler,
+            final Consumer<StreamMediator<ObjectEvent<BlockItem>, BlockItem>> shutdownCallback) {
+        this(new HashMap<>(), blockPersistenceHandler, shutdownCallback);
+    }
+
     @Override
     public void publishEvent(BlockItem blockItem) {
 
@@ -101,6 +112,7 @@ public void subscribe(final BlockItemEventHandler<ObjectEvent<BlockItem>> handle
         ringBuffer.addGatingSequences(batchEventProcessor.getSequence());
         executor.execute(batchEventProcessor);
 
+        // Keep track of the subscriber
         subscribers.put(handler, batchEventProcessor);
     }