1 files changed, 37 insertions, 18 deletions

diff --git a/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java b/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java
index 33aa4357..78f415c2 100644
--- a/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java
+++ b/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java
@@ -35,15 +35,13 @@ import protobuf_unittest.UnittestProto.SparseEnumMessage;
 import protobuf_unittest.UnittestProto.TestAllTypes;
 import protobuf_unittest.UnittestProto.TestPackedTypes;
 import protobuf_unittest.UnittestProto.TestSparseEnum;
-
-import junit.framework.TestCase;
-
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import junit.framework.TestCase;
 
 /**
  * Unit test for {@link CodedOutputStream}.
@@ -151,16 +149,21 @@ public class CodedOutputStreamTest extends TestCase {
     private final int initialPosition;
     private final CodedOutputStream stream;
     private final ByteBuffer buffer;
+    private final boolean unsafe;
 
-    NioDirectCoder(int size) {
-      this(size, 0);
+    NioDirectCoder(int size, boolean unsafe) {
+      this(size, 0, unsafe);
     }
 
-    NioDirectCoder(int size, int initialPosition) {
+    NioDirectCoder(int size, int initialPosition, boolean unsafe) {
+      this.unsafe = unsafe;
       this.initialPosition = initialPosition;
       buffer = ByteBuffer.allocateDirect(size);
       buffer.position(initialPosition);
-      stream = CodedOutputStream.newInstance(buffer);
+      stream =
+          unsafe
+              ? CodedOutputStream.newUnsafeInstance(buffer)
+              : CodedOutputStream.newSafeInstance(buffer);
     }
 
     @Override
@@ -181,7 +184,7 @@ public class CodedOutputStreamTest extends TestCase {
 
     @Override
     public OutputType getOutputType() {
-      return OutputType.NIO_DIRECT;
+      return unsafe ? OutputType.NIO_DIRECT_SAFE : OutputType.NIO_DIRECT_UNSAFE;
     }
   }
 
@@ -198,10 +201,16 @@ public class CodedOutputStreamTest extends TestCase {
         return new NioHeapCoder(size);
       }
     },
-    NIO_DIRECT() {
+    NIO_DIRECT_SAFE() {
       @Override
       Coder newCoder(int size) {
-        return new NioDirectCoder(size);
+        return new NioDirectCoder(size, false);
+      }
+    },
+    NIO_DIRECT_UNSAFE() {
+      @Override
+      Coder newCoder(int size) {
+        return new NioDirectCoder(size, true);
       }
     },
     STREAM() {
@@ -389,6 +398,7 @@ public class CodedOutputStreamTest extends TestCase {
         != CodedOutputStream.computeUInt32SizeNoTag(string.length() * Utf8.MAX_BYTES_PER_CHAR));
 
     coder.stream().writeStringNoTag(string);
+    coder.stream().flush();
     int stringSize = CodedOutputStream.computeStringSizeNoTag(string);
 
     // Verify that the total bytes written is correct
@@ -478,11 +488,12 @@ public class CodedOutputStreamTest extends TestCase {
 
   public void testWriteByteArrayWithOffsets() throws Exception {
     byte[] fullArray = bytes(0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88);
-    byte[] destination = new byte[4];
-    CodedOutputStream codedStream = CodedOutputStream.newInstance(destination);
-    codedStream.writeByteArrayNoTag(fullArray, 2, 2);
-    assertEqualBytes(OutputType.ARRAY, bytes(0x02, 0x33, 0x44, 0x00), destination);
-    assertEquals(3, codedStream.getTotalBytesWritten());
+    for (OutputType type : new OutputType[] {OutputType.ARRAY}) {
+      Coder coder = type.newCoder(4);
+      coder.stream().writeByteArrayNoTag(fullArray, 2, 2);
+      assertEqualBytes(type, bytes(0x02, 0x33, 0x44), coder.toByteArray());
+      assertEquals(3, coder.stream().getTotalBytesWritten());
+    }
   }
 
   public void testSerializeUtf8_MultipleSmallWrites() throws Exception {
@@ -561,7 +572,12 @@ public class CodedOutputStreamTest extends TestCase {
     // Tag is one byte, varint describing string length is 1 byte, string length is 9 bytes.
     // An array of size 1 will cause a failure when trying to write the varint.
     for (OutputType outputType :
-        new OutputType[] {OutputType.ARRAY, OutputType.NIO_HEAP, OutputType.NIO_DIRECT}) {
+        new OutputType[] {
+          OutputType.ARRAY,
+          OutputType.NIO_HEAP,
+          OutputType.NIO_DIRECT_SAFE,
+          OutputType.NIO_DIRECT_UNSAFE
+        }) {
       for (int i = 0; i < 11; i++) {
         Coder coder = outputType.newCoder(i);
         try {
@@ -599,10 +615,13 @@ public class CodedOutputStreamTest extends TestCase {
 
   public void testNioEncodersWithInitialOffsets() throws Exception {
     String value = "abc";
-    for (Coder coder : new Coder[] {new NioHeapCoder(10, 2), new NioDirectCoder(10, 2)}) {
+    for (Coder coder :
+        new Coder[] {
+          new NioHeapCoder(10, 2), new NioDirectCoder(10, 2, false), new NioDirectCoder(10, 2, true)
+        }) {
       coder.stream().writeStringNoTag(value);
       coder.stream().flush();
-      assertEqualBytes(coder.getOutputType(), new byte[]{3, 'a', 'b', 'c'}, coder.toByteArray());
+      assertEqualBytes(coder.getOutputType(), new byte[] {3, 'a', 'b', 'c'}, coder.toByteArray());
     }
   }