datafusion-datasource-json 53.1.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Utility types for JSON processing

use std::io::{BufRead, Read};

use bytes::Bytes;

// ============================================================================
// JsonArrayToNdjsonReader - Streaming JSON Array to NDJSON Converter
// ============================================================================
//
// Architecture:
//
// ```text
// ┌─────────────────────────────────────────────────────────────┐
// │  JSON Array File (potentially very large, e.g. 33GB)       │
// │  [{"a":1}, {"a":2}, {"a":3}, ...... {"a":1000000}]         │
// └─────────────────────────────────────────────────────────────┘
//                           │
//                           ▼ read chunks via ChannelReader
//                 ┌───────────────────┐
//                 │ JsonArrayToNdjson │  ← character substitution only:
//                 │      Reader       │    '[' skip, ',' → '\n', ']' stop
//                 └───────────────────┘
//                           │
//                           ▼ outputs NDJSON format
//                 ┌───────────────────┐
//                 │   Arrow Reader    │  ← internal buffer, batch parsing
//                 │  batch_size=8192  │
//                 └───────────────────┘
//                           │
//                           ▼ outputs RecordBatch
//                 ┌───────────────────┐
//                 │   RecordBatch     │
//                 └───────────────────┘
// ```
//
// Memory Efficiency:
//
// | Approach                              | Memory for 33GB file | Parse count |
// |---------------------------------------|----------------------|-------------|
// | Load entire file + serde_json         | ~100GB+              | 3x          |
// | Streaming with JsonArrayToNdjsonReader| ~32MB (configurable) | 1x          |
//
// Design Note:
//
// This implementation uses `inner: R` directly (not `BufReader<R>`) and manages
// its own input buffer. This is critical for compatibility with `SyncIoBridge`
// and `ChannelReader` in `spawn_blocking` contexts.
//

/// Default buffer size for JsonArrayToNdjsonReader (2MB for better throughput)
const DEFAULT_BUF_SIZE: usize = 2 * 1024 * 1024;

/// Parser state for JSON array streaming
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum JsonArrayState {
    /// Initial state, looking for opening '['
    Start,
    /// Inside the JSON array, processing objects
    InArray,
    /// Reached the closing ']', finished
    Done,
}

/// A streaming reader that converts JSON array format to NDJSON format.
///
/// This reader wraps an underlying reader containing JSON array data
/// `[{...}, {...}, ...]` and transforms it on-the-fly to newline-delimited
/// JSON format that Arrow's JSON reader can process.
///
/// Implements both `Read` and `BufRead` traits for compatibility with Arrow's
/// `ReaderBuilder::build()` which requires `BufRead`.
///
/// # Transformation Rules
///
/// - Skip leading `[` and whitespace before it
/// - Convert top-level `,` (between objects) to `\n`
/// - Skip whitespace at top level (between objects)
/// - Stop at trailing `]`
/// - Preserve everything inside objects (including nested `[`, `]`, `,`)
/// - Properly handle strings (ignore special chars inside quotes)
///
/// # Example
///
/// ```text
/// Input:  [{"a":1}, {"b":[1,2]}, {"c":"x,y"}]
/// Output: {"a":1}
///         {"b":[1,2]}
///         {"c":"x,y"}
/// ```
pub struct JsonArrayToNdjsonReader<R: Read> {
    /// Inner reader - we use R directly (not `BufReader<R>`) for SyncIoBridge compatibility
    inner: R,
    state: JsonArrayState,
    /// Tracks nesting depth of `{` and `[` to identify top-level commas
    depth: i32,
    /// Whether we're currently inside a JSON string
    in_string: bool,
    /// Whether the next character is escaped (after `\`)
    escape_next: bool,
    /// Input buffer - stores raw bytes read from inner reader
    input_buffer: Vec<u8>,
    /// Current read position in input buffer
    input_pos: usize,
    /// Number of valid bytes in input buffer
    input_filled: usize,
    /// Output buffer - stores transformed NDJSON bytes
    output_buffer: Vec<u8>,
    /// Current read position in output buffer
    output_pos: usize,
    /// Number of valid bytes in output buffer
    output_filled: usize,
    /// Whether trailing non-whitespace content was detected after ']'
    has_trailing_content: bool,
    /// Whether leading non-whitespace content was detected before '['
    has_leading_content: bool,
}

impl<R: Read> JsonArrayToNdjsonReader<R> {
    /// Create a new streaming reader that converts JSON array to NDJSON.
    pub fn new(reader: R) -> Self {
        Self::with_capacity(reader, DEFAULT_BUF_SIZE)
    }

    /// Create a new streaming reader with custom buffer size.
    ///
    /// Larger buffers improve throughput but use more memory.
    /// Total memory usage is approximately 2 * capacity (input + output buffers).
    pub fn with_capacity(reader: R, capacity: usize) -> Self {
        Self {
            inner: reader,
            state: JsonArrayState::Start,
            depth: 0,
            in_string: false,
            escape_next: false,
            input_buffer: vec![0; capacity],
            input_pos: 0,
            input_filled: 0,
            output_buffer: vec![0; capacity],
            output_pos: 0,
            output_filled: 0,
            has_trailing_content: false,
            has_leading_content: false,
        }
    }

    /// Check if the JSON array was properly terminated.
    ///
    /// This should be called after all data has been read.
    ///
    /// Returns an error if:
    /// - Unbalanced braces/brackets (depth != 0)
    /// - Unterminated string
    /// - Missing closing `]`
    /// - Unexpected trailing content after `]`
    pub fn validate_complete(&self) -> std::io::Result<()> {
        if self.has_leading_content {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Malformed JSON: unexpected leading content before '['",
            ));
        }
        if self.depth != 0 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Malformed JSON array: unbalanced braces or brackets",
            ));
        }
        if self.in_string {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Malformed JSON array: unterminated string",
            ));
        }
        if self.state != JsonArrayState::Done {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Incomplete JSON array: expected closing bracket ']'",
            ));
        }
        if self.has_trailing_content {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Malformed JSON: unexpected trailing content after ']'",
            ));
        }
        Ok(())
    }

    /// Process a single byte and return the transformed byte (if any)
    #[inline]
    fn process_byte(&mut self, byte: u8) -> Option<u8> {
        match self.state {
            JsonArrayState::Start => {
                // Looking for the opening '[', skip whitespace
                if byte == b'[' {
                    self.state = JsonArrayState::InArray;
                } else if !byte.is_ascii_whitespace() {
                    self.has_leading_content = true;
                }
                None
            }
            JsonArrayState::InArray => {
                // Handle escape sequences in strings
                if self.escape_next {
                    self.escape_next = false;
                    return Some(byte);
                }

                if self.in_string {
                    // Inside a string: handle escape and closing quote
                    match byte {
                        b'\\' => self.escape_next = true,
                        b'"' => self.in_string = false,
                        _ => {}
                    }
                    Some(byte)
                } else {
                    // Outside strings: track depth and transform
                    match byte {
                        b'"' => {
                            self.in_string = true;
                            Some(byte)
                        }
                        b'{' | b'[' => {
                            self.depth += 1;
                            Some(byte)
                        }
                        b'}' => {
                            self.depth -= 1;
                            Some(byte)
                        }
                        b']' => {
                            if self.depth == 0 {
                                // Top-level ']' means end of array
                                self.state = JsonArrayState::Done;
                                None
                            } else {
                                // Nested ']' inside an object
                                self.depth -= 1;
                                Some(byte)
                            }
                        }
                        b',' if self.depth == 0 => {
                            // Top-level comma between objects → newline
                            Some(b'\n')
                        }
                        _ => {
                            // At depth 0, skip whitespace between objects
                            if self.depth == 0 && byte.is_ascii_whitespace() {
                                None
                            } else {
                                Some(byte)
                            }
                        }
                    }
                }
            }
            JsonArrayState::Done => {
                // After ']', check for non-whitespace trailing content
                if !byte.is_ascii_whitespace() {
                    self.has_trailing_content = true;
                }
                None
            }
        }
    }

    /// Refill input buffer from inner reader if needed.
    /// Returns true if there's data available, false on EOF.
    fn refill_input_if_needed(&mut self) -> std::io::Result<bool> {
        if self.input_pos >= self.input_filled {
            // Input buffer exhausted, read more from inner
            let bytes_read = self.inner.read(&mut self.input_buffer)?;
            if bytes_read == 0 {
                return Ok(false); // EOF
            }
            self.input_pos = 0;
            self.input_filled = bytes_read;
        }
        Ok(true)
    }

    /// Fill the output buffer with transformed data.
    ///
    /// This method manages its own input buffer, reading from the inner reader
    /// as needed. When the output buffer is full, we stop processing but preserve
    /// the current position in the input buffer for the next call.
    fn fill_output_buffer(&mut self) -> std::io::Result<()> {
        let mut write_pos = 0;

        while write_pos < self.output_buffer.len() {
            // Refill input buffer if exhausted
            if !self.refill_input_if_needed()? {
                break; // EOF
            }

            // Process bytes from input buffer
            while self.input_pos < self.input_filled
                && write_pos < self.output_buffer.len()
            {
                let byte = self.input_buffer[self.input_pos];
                self.input_pos += 1;

                if let Some(transformed) = self.process_byte(byte) {
                    self.output_buffer[write_pos] = transformed;
                    write_pos += 1;
                }
            }
        }

        self.output_pos = 0;
        self.output_filled = write_pos;
        Ok(())
    }
}

impl<R: Read> Read for JsonArrayToNdjsonReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        // If output buffer is empty, fill it
        if self.output_pos >= self.output_filled {
            self.fill_output_buffer()?;
            if self.output_filled == 0 {
                return Ok(0); // EOF
            }
        }

        // Copy from output buffer to caller's buffer
        let available = self.output_filled - self.output_pos;
        let to_copy = std::cmp::min(available, buf.len());
        buf[..to_copy].copy_from_slice(
            &self.output_buffer[self.output_pos..self.output_pos + to_copy],
        );
        self.output_pos += to_copy;
        Ok(to_copy)
    }
}

impl<R: Read> BufRead for JsonArrayToNdjsonReader<R> {
    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
        if self.output_pos >= self.output_filled {
            self.fill_output_buffer()?;
        }
        Ok(&self.output_buffer[self.output_pos..self.output_filled])
    }

    fn consume(&mut self, amt: usize) {
        self.output_pos = std::cmp::min(self.output_pos + amt, self.output_filled);
    }
}

// ============================================================================
// ChannelReader - Sync reader that receives bytes from async channel
// ============================================================================
//
// Architecture:
//
// ```text
// ┌─────────────────────────────────────────────────────────────────────────┐
// │                         S3 / MinIO (async)                              │
// │                    (33GB JSON Array File)                               │
// └─────────────────────────────────────────────────────────────────────────┘
//                                 │
//                                 ▼ async stream (Bytes chunks)
// ┌─────────────────────────────────────────────────────────────────────────┐
// │                      Async Task (tokio runtime)                         │
// │              while let Some(chunk) = stream.next().await                │
// │                     byte_tx.send(chunk)                                 │
// └─────────────────────────────────────────────────────────────────────────┘
//                                 │
//                                 ▼ tokio::sync::mpsc::channel<Bytes>
//                                 │   (bounded, ~32MB buffer)
//                                 ▼
// ┌─────────────────────────────────────────────────────────────────────────┐
// │                   Blocking Task (spawn_blocking)                        │
// │  ┌──────────────┐   ┌────────────────────────┐   ┌──────────────────┐  │
// │  │ChannelReader │ → │JsonArrayToNdjsonReader │ → │ Arrow JsonReader │  │
// │  │   (Read)     │   │  [{},...] → {}\n{}     │   │  (RecordBatch)   │  │
// │  └──────────────┘   └────────────────────────┘   └──────────────────┘  │
// └─────────────────────────────────────────────────────────────────────────┘
//                                 │
//                                 ▼ tokio::sync::mpsc::channel<RecordBatch>
// ┌─────────────────────────────────────────────────────────────────────────┐
// │                      ReceiverStream (async)                             │
// │                   → DataFusion execution engine                         │
// └─────────────────────────────────────────────────────────────────────────┘
// ```
//
// Memory Budget (~32MB total):
// - sync_channel buffer: 128 chunks × ~128KB = ~16MB
// - JsonArrayToNdjsonReader: 2 × 2MB = 4MB
// - Arrow JsonReader internal: ~8MB
// - Miscellaneous: ~4MB
//

/// A synchronous `Read` implementation that receives bytes from an async channel.
///
/// This enables true streaming between async and sync contexts without
/// loading the entire file into memory. Uses `tokio::sync::mpsc::Receiver`
/// with `blocking_recv()` so the async producer never blocks a tokio worker
/// thread, while the sync consumer (running in `spawn_blocking`) safely blocks.
pub struct ChannelReader {
    rx: tokio::sync::mpsc::Receiver<Bytes>,
    current: Option<Bytes>,
    pos: usize,
}

impl ChannelReader {
    /// Create a new ChannelReader from a tokio mpsc receiver.
    pub fn new(rx: tokio::sync::mpsc::Receiver<Bytes>) -> Self {
        Self {
            rx,
            current: None,
            pos: 0,
        }
    }
}

impl Read for ChannelReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        loop {
            // If we have current chunk with remaining data, read from it
            if let Some(ref chunk) = self.current {
                let remaining = chunk.len() - self.pos;
                if remaining > 0 {
                    let to_copy = std::cmp::min(remaining, buf.len());
                    buf[..to_copy].copy_from_slice(&chunk[self.pos..self.pos + to_copy]);
                    self.pos += to_copy;
                    return Ok(to_copy);
                }
            }

            // Current chunk exhausted, get next from channel
            match self.rx.blocking_recv() {
                Some(bytes) => {
                    self.current = Some(bytes);
                    self.pos = 0;
                    // Loop back to read from new chunk
                }
                None => return Ok(0), // Channel closed = EOF
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_json_array_to_ndjson_simple() {
        let input = r#"[{"a":1}, {"a":2}, {"a":3}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        assert_eq!(output, "{\"a\":1}\n{\"a\":2}\n{\"a\":3}");
    }

    #[test]
    fn test_json_array_to_ndjson_nested() {
        let input = r#"[{"a":{"b":1}}, {"c":[1,2,3]}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        assert_eq!(output, "{\"a\":{\"b\":1}}\n{\"c\":[1,2,3]}");
    }

    #[test]
    fn test_json_array_to_ndjson_strings_with_special_chars() {
        let input = r#"[{"a":"[1,2]"}, {"b":"x,y"}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        assert_eq!(output, "{\"a\":\"[1,2]\"}\n{\"b\":\"x,y\"}");
    }

    #[test]
    fn test_json_array_to_ndjson_escaped_quotes() {
        let input = r#"[{"a":"say \"hello\""}, {"b":1}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        assert_eq!(output, "{\"a\":\"say \\\"hello\\\"\"}\n{\"b\":1}");
    }

    #[test]
    fn test_json_array_to_ndjson_empty() {
        let input = r#"[]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        assert_eq!(output, "");
    }

    #[test]
    fn test_json_array_to_ndjson_single_element() {
        let input = r#"[{"a":1}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        assert_eq!(output, "{\"a\":1}");
    }

    #[test]
    fn test_json_array_to_ndjson_bufread() {
        let input = r#"[{"a":1}, {"a":2}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());

        let buf = reader.fill_buf().unwrap();
        assert!(!buf.is_empty());

        let first_len = buf.len();
        reader.consume(first_len);

        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
    }

    #[test]
    fn test_json_array_to_ndjson_whitespace() {
        let input = r#"  [  {"a":1}  ,  {"a":2}  ]  "#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        // Top-level whitespace is skipped, internal whitespace preserved
        assert_eq!(output, "{\"a\":1}\n{\"a\":2}");
    }

    #[test]
    fn test_validate_complete_valid_json() {
        let valid_json = r#"[{"a":1},{"a":2}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(valid_json.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        reader.validate_complete().unwrap();
    }

    #[test]
    fn test_json_array_with_trailing_junk() {
        let input = r#" [ {"a":1} , {"a":2} ] some { junk [ here ] "#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        // Should extract the valid array content
        assert_eq!(output, "{\"a\":1}\n{\"a\":2}");

        // But validation should catch the trailing junk
        let result = reader.validate_complete();
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(
            err_msg.contains("trailing content")
                || err_msg.contains("Unexpected trailing"),
            "Expected trailing content error, got: {err_msg}"
        );
    }

    #[test]
    fn test_validate_complete_incomplete_array() {
        let invalid_json = r#"[{"a":1},{"a":2}"#; // Missing closing ]
        let mut reader = JsonArrayToNdjsonReader::new(invalid_json.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        let result = reader.validate_complete();
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(
            err_msg.contains("expected closing bracket")
                || err_msg.contains("missing closing"),
            "Expected missing bracket error, got: {err_msg}"
        );
    }

    #[test]
    fn test_validate_complete_unbalanced_braces() {
        let invalid_json = r#"[{"a":1},{"a":2]"#; // Wrong closing bracket
        let mut reader = JsonArrayToNdjsonReader::new(invalid_json.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        let result = reader.validate_complete();
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(
            err_msg.contains("unbalanced")
                || err_msg.contains("expected closing bracket"),
            "Expected unbalanced or missing bracket error, got: {err_msg}"
        );
    }

    #[test]
    fn test_json_array_with_leading_junk() {
        let input = r#"junk[{"a":1}, {"a":2}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        // Should still extract the valid array content
        assert_eq!(output, "{\"a\":1}\n{\"a\":2}");

        // But validation should catch the leading junk
        let result = reader.validate_complete();
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(
            err_msg.contains("leading content"),
            "Expected leading content error, got: {err_msg}"
        );
    }

    #[test]
    fn test_json_array_with_leading_whitespace_ok() {
        let input = r#"
  [{"a":1}, {"a":2}]"#;
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();
        assert_eq!(output, "{\"a\":1}\n{\"a\":2}");

        // Leading whitespace should be fine
        reader.validate_complete().unwrap();
    }

    #[test]
    fn test_validate_complete_valid_with_trailing_whitespace() {
        let input = r#"[{"a":1},{"a":2}]
    "#; // Trailing whitespace is OK
        let mut reader = JsonArrayToNdjsonReader::new(input.as_bytes());
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        // Whitespace after ] should be allowed
        reader.validate_complete().unwrap();
    }

    /// Test that data is not lost at buffer boundaries.
    ///
    /// This test creates input larger than the internal buffer to verify
    /// that newline characters are not dropped when they occur at buffer boundaries.
    #[test]
    fn test_buffer_boundary_no_data_loss() {
        // Create objects ~9KB each, so 10 objects = ~90KB
        let large_value = "x".repeat(9000);

        let mut objects = vec![];
        for i in 0..10 {
            objects.push(format!(r#"{{"id":{i},"data":"{large_value}"}}"#));
        }

        let input = format!("[{}]", objects.join(","));

        // Use small buffer to force multiple fill cycles
        let mut reader = JsonArrayToNdjsonReader::with_capacity(input.as_bytes(), 8192);
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        // Verify correct number of newlines (9 newlines separate 10 objects)
        let newline_count = output.matches('\n').count();
        assert_eq!(
            newline_count, 9,
            "Expected 9 newlines separating 10 objects, got {newline_count}"
        );

        // Verify each line is valid JSON
        for (i, line) in output.lines().enumerate() {
            let parsed: Result<serde_json::Value, _> = serde_json::from_str(line);
            assert!(
                parsed.is_ok(),
                "Line {} is not valid JSON: {}...",
                i,
                &line[..100.min(line.len())]
            );

            // Verify the id field matches expected value
            let value = parsed.unwrap();
            assert_eq!(
                value["id"].as_i64(),
                Some(i as i64),
                "Object {i} has wrong id"
            );
        }
    }

    /// Test with real-world-like data format (with leading whitespace and newlines)
    #[test]
    fn test_real_world_format_large() {
        let large_value = "x".repeat(8000);

        // Format similar to real files: opening bracket on its own line,
        // each object indented with 2 spaces
        let mut objects = vec![];
        for i in 0..10 {
            objects.push(format!(r#"  {{"id":{i},"data":"{large_value}"}}"#));
        }

        let input = format!("[\n{}\n]", objects.join(",\n"));

        let mut reader = JsonArrayToNdjsonReader::with_capacity(input.as_bytes(), 8192);
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        let lines: Vec<&str> = output.lines().collect();
        assert_eq!(lines.len(), 10, "Expected 10 objects");

        for (i, line) in lines.iter().enumerate() {
            assert!(
                line.starts_with("{\"id\""),
                "Line {} should start with object, got: {}...",
                i,
                &line[..50.min(line.len())]
            );
        }
    }

    /// Test ChannelReader
    #[test]
    fn test_channel_reader() {
        let (tx, rx) = tokio::sync::mpsc::channel(4);

        // Send some chunks (try_send is non-async)
        tx.try_send(Bytes::from("Hello, ")).unwrap();
        tx.try_send(Bytes::from("World!")).unwrap();
        drop(tx); // Close channel

        let mut reader = ChannelReader::new(rx);
        let mut output = String::new();
        reader.read_to_string(&mut output).unwrap();

        assert_eq!(output, "Hello, World!");
    }

    /// Test ChannelReader with small reads
    #[test]
    fn test_channel_reader_small_reads() {
        let (tx, rx) = tokio::sync::mpsc::channel(4);

        tx.try_send(Bytes::from("ABCDEFGHIJ")).unwrap();
        drop(tx);

        let mut reader = ChannelReader::new(rx);
        let mut buf = [0u8; 3];

        // Read in small chunks
        assert_eq!(reader.read(&mut buf).unwrap(), 3);
        assert_eq!(&buf, b"ABC");

        assert_eq!(reader.read(&mut buf).unwrap(), 3);
        assert_eq!(&buf, b"DEF");

        assert_eq!(reader.read(&mut buf).unwrap(), 3);
        assert_eq!(&buf, b"GHI");

        assert_eq!(reader.read(&mut buf).unwrap(), 1);
        assert_eq!(&buf[..1], b"J");

        // EOF
        assert_eq!(reader.read(&mut buf).unwrap(), 0);
    }
}