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 index 844e144ffc0..b9dea7d3f5c 100644 --- 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; @@ -250,18 +249,15 @@ void sendRequest_interruptBeforeInitialPayloadWritten( @SuppressWarnings("UnnecessaryAsync") private Class executeInterrupts( final boolean closeStream, final boolean exceedTimeout) { - final AtomicReference> expectedException = - new AtomicReference<>(null); + Class 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 getInterruptParams() {