diff --git a/test/framework/src/test/java/org/elasticsearch/transport/nio/WriteOperationTests.java b/test/framework/src/test/java/org/elasticsearch/transport/nio/WriteOperationTests.java index 0015d39a373..0085a9b204a 100644 --- a/test/framework/src/test/java/org/elasticsearch/transport/nio/WriteOperationTests.java +++ b/test/framework/src/test/java/org/elasticsearch/transport/nio/WriteOperationTests.java @@ -21,12 +21,15 @@ package org.elasticsearch.transport.nio; import org.elasticsearch.action.ActionListener; import org.elasticsearch.common.bytes.BytesArray; +import org.elasticsearch.common.bytes.CompositeBytesReference; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.transport.nio.channel.NioSocketChannel; import org.junit.Before; +import org.mockito.ArgumentCaptor; import java.io.IOException; import java.nio.ByteBuffer; +import java.util.List; import static org.mockito.Matchers.any; import static org.mockito.Mockito.mock; @@ -65,4 +68,52 @@ public class WriteOperationTests extends ESTestCase { assertFalse(writeOp.isFullyFlushed()); } + + public void testMultipleFlushesWithCompositeBuffer() throws IOException { + BytesArray bytesReference1 = new BytesArray(new byte[10]); + BytesArray bytesReference2 = new BytesArray(new byte[15]); + BytesArray bytesReference3 = new BytesArray(new byte[3]); + CompositeBytesReference bytesReference = new CompositeBytesReference(bytesReference1, bytesReference2, bytesReference3); + WriteOperation writeOp = new WriteOperation(channel, bytesReference, listener); + + ArgumentCaptor buffersCaptor = ArgumentCaptor.forClass(ByteBuffer[].class); + + when(channel.write(buffersCaptor.capture())).thenReturn(5) + .thenReturn(5) + .thenReturn(2) + .thenReturn(15) + .thenReturn(1); + + writeOp.flush(); + assertFalse(writeOp.isFullyFlushed()); + writeOp.flush(); + assertFalse(writeOp.isFullyFlushed()); + writeOp.flush(); + assertFalse(writeOp.isFullyFlushed()); + writeOp.flush(); + assertFalse(writeOp.isFullyFlushed()); + writeOp.flush(); + assertTrue(writeOp.isFullyFlushed()); + + List values = buffersCaptor.getAllValues(); + ByteBuffer[] byteBuffers = values.get(0); + assertEquals(3, byteBuffers.length); + assertEquals(10, byteBuffers[0].remaining()); + + byteBuffers = values.get(1); + assertEquals(3, byteBuffers.length); + assertEquals(5, byteBuffers[0].remaining()); + + byteBuffers = values.get(2); + assertEquals(2, byteBuffers.length); + assertEquals(15, byteBuffers[0].remaining()); + + byteBuffers = values.get(3); + assertEquals(2, byteBuffers.length); + assertEquals(13, byteBuffers[0].remaining()); + + byteBuffers = values.get(4); + assertEquals(1, byteBuffers.length); + assertEquals(1, byteBuffers[0].remaining()); + } }