PartitionRequestClientHandler extends ChannelInboundHandlerAdapter { void addInputChannel(RemoteInputChannel listener) throws IOException { checkError(); if (!inputChannels.containsKey(listener.getInputChannelId())) { inputChannels.put(listener.getInputChannelId(), listener); } }  @Override void channelActive(final ChannelHandlerContext ctx); @Override void channelInactive(ChannelHandlerContext ctx); @Override void exceptionCaught(ChannelHandlerContext ctx, Throwable cause); @Override void channelRead(ChannelHandlerContext ctx, Object msg); @Override void channelReadComplete(ChannelHandlerContext ctx);  }
@Test @SuppressWarnings("unchecked") public void testAutoReadAfterUnsuccessfulStagedMessage() throws Exception { PartitionRequestClientHandler handler = new PartitionRequestClientHandler(); EmbeddedChannel channel = new EmbeddedChannel(handler); final AtomicReference<BufferListener> listener = new AtomicReference<>(); BufferProvider bufferProvider = mock(BufferProvider.class); when(bufferProvider.addBufferListener(any(BufferListener.class))).thenAnswer(new Answer<Boolean>() { @Override @SuppressWarnings("unchecked") public Boolean answer(InvocationOnMock invocation) throws Throwable { listener.set((BufferListener) invocation.getArguments()[0]); return true; } }); when(bufferProvider.requestBuffer()).thenReturn(null); InputChannelID channelId = new InputChannelID(0, 0); RemoteInputChannel inputChannel = mock(RemoteInputChannel.class); when(inputChannel.getInputChannelId()).thenReturn(channelId); when(inputChannel.getBufferProvider()).thenReturn(bufferProvider, bufferProvider, null); handler.addInputChannel(inputChannel); BufferResponse msg = createBufferResponse(createBuffer(true), 0, channelId); assertTrue(channel.config().isAutoRead()); channel.writeInbound(msg); assertFalse(channel.config().isAutoRead()); msg = createBufferResponse(createBuffer(true), 1, channelId); channel.writeInbound(msg); msg = createBufferResponse(createBuffer(true), 2, channelId); channel.writeInbound(msg); Buffer availableBuffer = createBuffer(false); listener.get().notifyBufferAvailable(availableBuffer); when(bufferProvider.addBufferListener(any(BufferListener.class))).thenReturn(false); when(bufferProvider.isDestroyed()).thenReturn(true); channel.runPendingTasks(); assertTrue(channel.config().isAutoRead()); }