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 909c2afa67f..844e144ffc0 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,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 executeInterrupts( final boolean closeStream, final boolean exceedTimeout) { - Class expectedException = null; + final AtomicReference> 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 getInterruptParams() {