fix assemble

StefanBratanov · Sep 24, 2023 · 5357697 · 5357697
1 parent 3087fc1
commit 5357697
Showing 1 changed file with 5 additions and 9 deletions.
diff --git a/networking/p2p/src/test/java/tech/pegasys/teku/networking/p2p/libp2p/rpc/RpcHandlerTest.java b/networking/p2p/src/test/java/tech/pegasys/teku/networking/p2p/libp2p/rpc/RpcHandlerTest.java
@@ -29,7 +29,6 @@
 import io.libp2p.core.multistream.ProtocolBinding;
 import io.libp2p.core.mux.StreamMuxer.Session;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicReference;
 import kotlin.Unit;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -249,18 +248,15 @@ void sendRequest_interruptBeforeInitialPayloadWritten(
 
   private Class<? extends Exception> executeInterrupts(
       final boolean closeStream, final boolean exceedTimeout) {
-    final AtomicReference<Class<? extends Exception>> expectedException =
-        new AtomicReference<>(null);
+    Class<? extends Exception> expectedException = null;
     if (closeStream) {
       closeFuture.complete(null);
-      expectedException.set(PeerDisconnectedException.class);
-    }
-    if (exceedTimeout) {
+      expectedException = PeerDisconnectedException.class;
+    } else if (exceedTimeout) {
       asyncRunner.executeQueuedActions();
-      expectedException.compareAndSet(null, StreamTimeoutException.class);
+      expectedException = StreamTimeoutException.class;
     }
-
-    return expectedException.get();
+    return expectedException;
   }
 
   public static java.util.stream.Stream<Arguments> getInterruptParams() {