revert changes

StefanBratanov · Oct 20, 2023 · b26a177 · b26a177
1 parent 4e93fdd
commit b26a177
Showing 1 changed file with 10 additions and 5 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,6 +29,7 @@
 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;
@@ -246,17 +247,21 @@ void sendRequest_interruptBeforeInitialPayloadWritten(
     verify(stream).close();
   }
 
+  @SuppressWarnings("UnnecessaryAsync")
   private Class<? extends Exception> executeInterrupts(
       final boolean closeStream, final boolean exceedTimeout) {
-    Class<? extends Exception> expectedException = null;
+    final AtomicReference<Class<? extends Exception>> expectedException =
+        new AtomicReference<>(null);
     if (closeStream) {
       closeFuture.complete(null);
-      expectedException = PeerDisconnectedException.class;
-    } else if (exceedTimeout) {
+      expectedException.set(PeerDisconnectedException.class);
+    }
+    if (exceedTimeout) {
       asyncRunner.executeQueuedActions();
-      expectedException = StreamTimeoutException.class;
+      expectedException.compareAndSet(null, StreamTimeoutException.class);
     }
-    return expectedException;
+
+    return expectedException.get();
   }
 
   public static java.util.stream.Stream<Arguments> getInterruptParams() {