merge tests

Author: xborder

flight/flight-core/pom.xml

@@ -146,6 +146,12 @@ under the License.
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-lang3</artifactId>
+      <version>3.20.0</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

…TestFlightDataParserDuplicateFields.java …/arrow/flight/TestArrowMessageParse.javaflight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightDataParserDuplicateFields.java renamed to flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageParse.java flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightDataParserDuplicateFields.java renamed to flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageParse.java

@@ -50,11 +50,11 @@
 import org.junit.jupiter.api.Test;
 
 /**
- * Tests FlightData parsing for duplicate field handling and well-formed messages. Covers both
- * InputStream (with copying) and ArrowBuf (zero-copy) parsing paths. Verifies that duplicate
- * protobuf fields use last-occurrence-wins semantics without memory leaks.
+ * Tests FlightData parsing including duplicate field handling, well-formed messages, and zero-copy
+ * behavior. Covers both InputStream (with copying) and ArrowBuf (zero-copy) parsing paths. Verifies
+ * that duplicate protobuf fields use last-occurrence-wins semantics without memory leaks.
  */
-public class TestFlightDataParserDuplicateFields {
+public class TestArrowMessageParse {
 
   private BufferAllocator allocator;
 
@@ -110,7 +110,7 @@ public void testDuplicateAppMetadataArrowBuf() throws Exception {
             List.of(
                 Pair.of(FlightData.APP_METADATA_FIELD_NUMBER, firstAppMetadata),
                 Pair.of(FlightData.APP_METADATA_FIELD_NUMBER, secondAppMetadata)));
-    InputStream stream = new DetachableDirectBufferInputStream(serialized);
+    InputStream stream = MockGrpcInputStream.ofDirectBuffer(serialized);
 
     try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) {
       ArrowBuf appMetadata = message.getApplicationMetadata();
@@ -166,7 +166,7 @@ public void testDuplicateBodyArrowBuf() throws Exception {
             List.of(
                 Pair.of(FlightData.DATA_BODY_FIELD_NUMBER, firstBody),
                 Pair.of(FlightData.DATA_BODY_FIELD_NUMBER, secondBody)));
-    InputStream stream = new DetachableDirectBufferInputStream(serialized);
+    InputStream stream = MockGrpcInputStream.ofDirectBuffer(serialized);
 
     try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) {
       ArrowBuf body = Iterables.getOnlyElement(message.getBufs());
@@ -229,7 +229,7 @@ public void testFieldsArrowBuf() throws Exception {
     assertEquals(0, allocator.getAllocatedMemory());
 
     byte[] serialized = buildFlightDataWithBothFields(appMetadataBytes, bodyBytes);
-    InputStream stream = new DetachableDirectBufferInputStream(serialized);
+    InputStream stream = MockGrpcInputStream.ofDirectBuffer(serialized);
 
     try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) {
       // Verify descriptor
@@ -260,6 +260,21 @@ public void testFieldsArrowBuf() throws Exception {
     assertEquals(0, allocator.getAllocatedMemory());
   }
 
+  /** Verifies that heap buffers fall back to InputStream path without calling detach(). */
+  @Test
+  public void testHeapBufferFallbackDoesNotDetach() throws Exception {
+    byte[] appMetadataBytes = new byte[] {8, 9};
+    byte[] bodyBytes = new byte[] {10, 11, 12};
+
+    byte[] serialized = buildFlightDataWithBothFields(appMetadataBytes, bodyBytes);
+    MockGrpcInputStream stream = MockGrpcInputStream.ofHeapBuffer(serialized);
+
+    try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) {
+      assertNotNull(message.getDescriptor());
+      assertEquals(0, stream.getDetachCount());
+    }
+  }
+
   // Helper methods to build complete FlightData messages
 
   private FlightDescriptor createTestDescriptor() {
@@ -298,7 +313,8 @@ private byte[] buildFlightDataWithBothFields(byte[] appMetadata, byte[] body) th
 
   // Helper methods to build FlightData messages with duplicate fields
 
-  private byte[] buildFlightDataDescriptors(List<Pair<Integer, byte[]>> descriptors) throws IOException {
+  private byte[] buildFlightDataDescriptors(List<Pair<Integer, byte[]>> descriptors)
+      throws IOException {
 
     ByteArrayOutputStream baos = new ByteArrayOutputStream();
     CodedOutputStream cos = CodedOutputStream.newInstance(baos);
@@ -311,34 +327,47 @@ private byte[] buildFlightDataDescriptors(List<Pair<Integer, byte[]>> descriptor
   }
 
   /** Mock InputStream implementing gRPC's Detachable and HasByteBuffer for testing zero-copy. */
-  private static class DetachableDirectBufferInputStream extends InputStream
+  private static class MockGrpcInputStream extends InputStream
       implements Detachable, HasByteBuffer {
     private ByteBuffer buffer;
+    private final boolean byteBufferSupported;
+    private int detachCount;
 
-    DetachableDirectBufferInputStream(byte[] data) {
-      this.buffer = ByteBuffer.allocateDirect(data.length);
-      this.buffer.put(data).flip();
+    private MockGrpcInputStream(ByteBuffer buffer, boolean byteBufferSupported) {
+      this.buffer = buffer;
+      this.byteBufferSupported = byteBufferSupported;
     }
 
-    private DetachableDirectBufferInputStream(ByteBuffer buffer) {
-      this.buffer = buffer;
+    static MockGrpcInputStream ofDirectBuffer(byte[] data) {
+      ByteBuffer buf = ByteBuffer.allocateDirect(data.length);
+      buf.put(data).flip();
+      return new MockGrpcInputStream(buf, true);
+    }
+
+    static MockGrpcInputStream ofHeapBuffer(byte[] data) {
+      return new MockGrpcInputStream(ByteBuffer.wrap(data), true);
     }
 
     @Override
     public boolean byteBufferSupported() {
-      return true;
+      return byteBufferSupported;
     }
 
     @Override
     public ByteBuffer getByteBuffer() {
-      return buffer;
+      return byteBufferSupported ? buffer : null;
     }
 
     @Override
     public InputStream detach() {
+      detachCount++;
       ByteBuffer detached = this.buffer;
       this.buffer = null;
-      return new DetachableDirectBufferInputStream(detached);
+      return new MockGrpcInputStream(detached, byteBufferSupported);
+    }
+
+    int getDetachCount() {
+      return detachCount;
     }
 
     @Override