5656import static org .elasticsearch .transport .TransportService .NOOP_TRANSPORT_INTERCEPTOR ;
5757import static org .hamcrest .Matchers .containsString ;
5858import static org .hamcrest .Matchers .endsWith ;
59+ import static org .hamcrest .Matchers .instanceOf ;
5960
6061public class DisruptableMockTransportTests extends ESTestCase {
6162
@@ -102,7 +103,7 @@ public void initTransports() {
102103 deterministicTaskQueue = new DeterministicTaskQueue (
103104 Settings .builder ().put (Node .NODE_NAME_SETTING .getKey (), "dummy" ).build (), random ());
104105
105- final DisruptableMockTransport transport1 = new DisruptableMockTransport (node1 , logger ) {
106+ final DisruptableMockTransport transport1 = new DisruptableMockTransport (node1 , logger , deterministicTaskQueue ) {
106107 @ Override
107108 protected ConnectionStatus getConnectionStatus (DiscoveryNode destination ) {
108109 return DisruptableMockTransportTests .this .getConnectionStatus (getLocalNode (), destination );
@@ -119,7 +120,7 @@ protected void execute(Runnable runnable) {
119120 }
120121 };
121122
122- final DisruptableMockTransport transport2 = new DisruptableMockTransport (node2 , logger ) {
123+ final DisruptableMockTransport transport2 = new DisruptableMockTransport (node2 , logger , deterministicTaskQueue ) {
123124 @ Override
124125 protected ConnectionStatus getConnectionStatus (DiscoveryNode destination ) {
125126 return DisruptableMockTransportTests .this .getConnectionStatus (getLocalNode (), destination );
@@ -317,27 +318,33 @@ public void testDisconnectedOnSuccessfulResponse() throws IOException {
317318 AtomicReference <TransportChannel > responseHandlerChannel = new AtomicReference <>();
318319 registerRequestHandler (service2 , requestHandlerCaptures (responseHandlerChannel ::set ));
319320
320- send (service1 , node2 , responseHandlerShouldNotBeCalled ());
321+ AtomicReference <TransportException > responseHandlerException = new AtomicReference <>();
322+ send (service1 , node2 , responseHandlerShouldBeCalledExceptionally (responseHandlerException ::set ));
321323 deterministicTaskQueue .runAllRunnableTasks ();
322324 assertNotNull (responseHandlerChannel .get ());
325+ assertNull (responseHandlerException .get ());
323326
324327 disconnectedLinks .add (Tuple .tuple (node2 , node1 ));
325328 responseHandlerChannel .get ().sendResponse (TransportResponse .Empty .INSTANCE );
326- deterministicTaskQueue .runAllRunnableTasks ();
329+ deterministicTaskQueue .runAllTasks ();
330+ assertThat (responseHandlerException .get (), instanceOf (ConnectTransportException .class ));
327331 }
328332
329333 public void testDisconnectedOnExceptionalResponse () throws IOException {
330334 registerRequestHandler (service1 , requestHandlerShouldNotBeCalled ());
331335 AtomicReference <TransportChannel > responseHandlerChannel = new AtomicReference <>();
332336 registerRequestHandler (service2 , requestHandlerCaptures (responseHandlerChannel ::set ));
333337
334- send (service1 , node2 , responseHandlerShouldNotBeCalled ());
338+ AtomicReference <TransportException > responseHandlerException = new AtomicReference <>();
339+ send (service1 , node2 , responseHandlerShouldBeCalledExceptionally (responseHandlerException ::set ));
335340 deterministicTaskQueue .runAllRunnableTasks ();
336341 assertNotNull (responseHandlerChannel .get ());
342+ assertNull (responseHandlerException .get ());
337343
338344 disconnectedLinks .add (Tuple .tuple (node2 , node1 ));
339345 responseHandlerChannel .get ().sendResponse (new Exception ());
340- deterministicTaskQueue .runAllRunnableTasks ();
346+ deterministicTaskQueue .runAllTasks ();
347+ assertThat (responseHandlerException .get (), instanceOf (ConnectTransportException .class ));
341348 }
342349
343350 public void testUnavailableOnSuccessfulResponse () throws IOException {
0 commit comments