Ported FragmentedMessageDeliveryTest to driver

Author: pontusmelke

driver/src/main/java/org/neo4j/driver/v1/util/Functions.java

@@ -24,6 +24,7 @@
  */
 public class Functions
 {
+    @SuppressWarnings( "unchecked" )
     public static <T> Function<T,T> identity()
     {
         return IDENTITY;

driver/src/test/java/org/neo4j/driver/internal/messaging/FragmentedMessageDeliveryTest.java

@@ -0,0 +1,165 @@
+/**
+ * Copyright (c) 2002-2016 "Neo Technology,"
+ * Network Engine for Objects in Lund AB [http://neotechnology.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.internal.messaging;
+
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
+import java.util.ArrayList;
+import java.util.Collections;
+
+import org.neo4j.driver.internal.connector.socket.ChunkedOutput;
+import org.neo4j.driver.v1.Value;
+import org.neo4j.driver.v1.util.DumpMessage;
+
+import static java.util.Arrays.asList;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+
+/**
+ * This tests network fragmentation of messages. Given a set of messages, it will serialize and chunk the message up
+ * to a specified chunk size. Then it will split that data into a specified number of fragments, trying every possible
+ * permutation of fragment sizes for the specified number. For instance, assuming an unfragmented message size of 15,
+ * and a fragment count of 3, it will create fragment size permutations like:
+ * <p>
+ * [1,1,13]
+ * [1,2,12]
+ * [1,3,11]
+ * ..
+ * [12,1,1]
+ * <p>
+ * For each permutation, it delivers the fragments to the protocol implementation, and asserts the protocol handled
+ * them properly.
+ */
+public class FragmentedMessageDeliveryTest
+{
+    private final MessageFormat format = new PackStreamMessageFormatV1();
+
+    // Only test one chunk size for now, this can be parameterized to test lots of different ones
+    private int chunkSize = 16;
+
+    // Only test one message for now. This can be parameterized later to test lots of different ones
+    private Message[] messages = new Message[]{ new RunMessage( "Mjölnir", Collections.<String, Value>emptyMap() )};
+
+    @Test
+    public void testFragmentedMessageDelivery() throws Throwable
+    {
+        // Given
+        byte[] unfragmented = serialize( messages );
+
+        // When & Then
+        int n = unfragmented.length;
+        for ( int i = 1; i < n - 1; i++ )
+        {
+            for ( int j = 1; j < n - i; j++ )
+            {
+                testPermutation( unfragmented, i, j, n - i - j );
+            }
+        }
+    }
+
+    private void testPermutation( byte[] unfragmented, int... sizes ) throws IOException
+    {
+        int pos = 0;
+        ByteBuffer[] fragments = new ByteBuffer[sizes.length];
+        for ( int i = 0; i < sizes.length; i++ )
+        {
+            fragments[i] = ByteBuffer.wrap( unfragmented, pos, sizes[i] );
+            pos += sizes[i];
+        }
+        testPermutation( unfragmented, fragments );
+    }
+
+    private void testPermutation( byte[] unfragmented, ByteBuffer[] fragments ) throws IOException
+    {
+
+        // When data arrives split up according to the current permutation
+        ReadableByteChannel[] channels = new ReadableByteChannel[fragments.length];
+        for ( int i = 0; i < fragments.length; i++ )
+        {
+            channels[i] = packet( fragments[i] );
+        }
+
+        ReadableByteChannel fragmentedChannel = packets( channels );
+        MessageFormat.Reader reader = format.newReader( fragmentedChannel );
+
+        ArrayList<Message> packedMessages = new ArrayList<>();
+        DumpMessage.unpack( packedMessages, reader );
+
+        assertThat( packedMessages, equalTo(asList(messages)) );
+    }
+
+    private ReadableByteChannel packet( ByteBuffer buffer )
+    {
+        //NOTE buffer.array is ok here since we know buffer is backed by array
+        return Channels.newChannel(
+                new ByteArrayInputStream( buffer.array() ) );
+    }
+
+
+    private ReadableByteChannel packets( final ReadableByteChannel... channels )
+    {
+
+        return new ReadableByteChannel()
+        {
+            private int index = 0;
+
+            @Override
+            public int read( ByteBuffer dst ) throws IOException
+            {
+                return channels[index++].read( dst );
+            }
+
+            @Override
+            public boolean isOpen()
+            {
+                return false;
+            }
+
+            @Override
+            public void close() throws IOException
+            {
+
+            }
+        };
+    }
+
+    private byte[] serialize( Message... msgs ) throws IOException
+    {
+
+            final ByteArrayOutputStream out = new ByteArrayOutputStream( 128 );
+
+        ChunkedOutput output = new ChunkedOutput(  chunkSize + 2 /* for chunk header */, Channels.newChannel( out ) );
+
+        PackStreamMessageFormatV1.Writer writer =
+                new PackStreamMessageFormatV1.Writer( output, output.messageBoundaryHook() );
+            for ( Message message : messages )
+            {
+                writer.write( message );
+            }
+            writer.flush();
+
+            return out.toByteArray();
+    }
+}

examples/src/test/java/org/neo4j/docs/driver/ExamplesIT.java

@@ -150,7 +150,7 @@ public void retainResultsForLaterProcessing() throws Throwable
     {
         StdIOCapture stdIO = new StdIOCapture();
         try ( AutoCloseable captured = stdIO.capture();
-                Driver driver = GraphDatabase.driver( "bolt://localhost" ); )
+                Driver driver = GraphDatabase.driver( "bolt://localhost" ) )
         {
             try ( Session setup = driver.session() )
             {
@@ -195,6 +195,7 @@ public void transactionRollback() throws Throwable
         }
     }
 
+    @SuppressWarnings( "unchecked" )
     @Test
     public void resultSummary() throws Throwable
     {