Add tests and fix BytesReader to match BytesIO.read() semantics

BytesReader.read(n) now returns partial data at end-of-buffer instead of
raising EOFError, matching BytesIO behavior. Added 11 unit tests for
BytesReader and 6 tests verifying _ConnectionIOBuffer reset position
semantics under both checksumming and non-checksumming paths.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: dkropachev

cassandra/protocol.py

@@ -57,8 +57,11 @@ class InternalError(Exception):
 class BytesReader:
     """
     Lightweight reader for bytes data without BytesIO overhead.
-    Provides the same read() interface but operates directly on a
-    bytes or memoryview object, avoiding internal buffer copies.
+    Provides the same read() interface as BytesIO but operates directly
+    on a bytes or memoryview object, avoiding internal buffer copies.
+
+    read(n) behaves like BytesIO.read(n): returns up to n bytes and
+    returns fewer bytes (or empty bytes) when the end of data is reached.
     """
     __slots__ = ('_data', '_pos', '_size')
 
@@ -72,9 +75,7 @@ def read(self, n=-1):
             result = self._data[self._pos:]
             self._pos = self._size
         else:
-            end = self._pos + n
-            if end > self._size:
-                raise EOFError("Cannot read past the end of the buffer")
+            end = min(self._pos + n, self._size)
             result = self._data[self._pos:end]
             self._pos = end
         # Return bytes to maintain compatibility with unpack functions

tests/unit/test_connection.py

@@ -22,7 +22,8 @@
 from cassandra.cluster import Cluster
 from cassandra.connection import (Connection, HEADER_DIRECTION_TO_CLIENT, ProtocolError,
                                   locally_supported_compressions, ConnectionHeartbeat, _Frame, Timer, TimerManager,
-                                  ConnectionException, ConnectionShutdown, DefaultEndPoint, ShardAwarePortGenerator)
+                                  ConnectionException, ConnectionShutdown, DefaultEndPoint, ShardAwarePortGenerator,
+                                  _ConnectionIOBuffer)
 from cassandra.marshal import uint8_pack, uint32_pack, int32_pack
 from cassandra.protocol import (write_stringmultimap, write_int, write_string,
                                 SupportedMessage, ProtocolHandler)
@@ -571,3 +572,80 @@ def test_generate_is_repeatable_with_same_mock(self, mock_randrange):
         second_run = list(itertools.islice(gen.generate(0, 2), 5))
 
         assert first_run == second_run
+
+
+class TestConnectionIOBufferReset(unittest.TestCase):
+    """Verify _reset_buffer and reset_cql_frame_buffer position semantics."""
+
+    def test_reset_buffer_discards_consumed_data(self):
+        buf = BytesIO(b'\x01\x02\x03\x04\x05')
+        buf.seek(0)
+        # Consume first 3 bytes
+        assert buf.read(3) == b'\x01\x02\x03'
+        new_buf = _ConnectionIOBuffer._reset_buffer(buf)
+        # New buffer should contain only unconsumed data
+        new_buf.seek(0)
+        assert new_buf.read() == b'\x04\x05'
+
+    def test_reset_buffer_position_at_end(self):
+        buf = BytesIO(b'\x01\x02\x03')
+        buf.seek(0)
+        buf.read(1)
+        new_buf = _ConnectionIOBuffer._reset_buffer(buf)
+        # Position should be at end (ready for appending)
+        assert new_buf.tell() == 2
+
+    def test_reset_buffer_fully_consumed(self):
+        buf = BytesIO(b'\x01\x02')
+        buf.seek(0)
+        buf.read(2)
+        new_buf = _ConnectionIOBuffer._reset_buffer(buf)
+        new_buf.seek(0)
+        assert new_buf.read() == b''
+
+    def test_reset_buffer_nothing_consumed(self):
+        buf = BytesIO(b'\x01\x02\x03')
+        buf.seek(0)
+        new_buf = _ConnectionIOBuffer._reset_buffer(buf)
+        new_buf.seek(0)
+        assert new_buf.read() == b'\x01\x02\x03'
+
+    @staticmethod
+    def _make_iobuf(checksumming=False):
+        conn = Mock()
+        conn._is_checksumming_enabled = checksumming
+        iobuf = _ConnectionIOBuffer(conn)
+        # Keep a strong reference so the weakref.proxy inside iobuf stays valid
+        iobuf._conn_ref = conn
+        if checksumming:
+            iobuf.set_checksumming_buffer()
+        return iobuf
+
+    def test_reset_cql_frame_buffer_checksumming_uses_tell_position(self):
+        """
+        When checksumming is enabled, reset_cql_frame_buffer delegates to
+        _reset_buffer which relies on tell() to determine consumed data.
+        Verify that seeking to an arbitrary position before reset correctly
+        preserves only the unconsumed tail.
+        """
+        iobuf = self._make_iobuf(checksumming=True)
+        # Write some data into the cql_frame_buffer
+        iobuf.cql_frame_buffer.write(b'\xAA\xBB\xCC\xDD\xEE')
+        # Seek to position 3, simulating that first 3 bytes were consumed
+        iobuf.cql_frame_buffer.seek(3)
+        iobuf.reset_cql_frame_buffer()
+        # After reset, only the unconsumed tail should remain
+        iobuf.cql_frame_buffer.seek(0)
+        assert iobuf.cql_frame_buffer.read() == b'\xDD\xEE'
+
+    def test_reset_cql_frame_buffer_no_checksumming_resets_io_buffer(self):
+        """
+        Without checksumming, reset_cql_frame_buffer delegates to
+        reset_io_buffer (since cql_frame_buffer IS the io_buffer).
+        """
+        iobuf = self._make_iobuf(checksumming=False)
+        iobuf.io_buffer.write(b'\x01\x02\x03\x04')
+        iobuf.io_buffer.seek(2)
+        iobuf.reset_cql_frame_buffer()
+        iobuf.io_buffer.seek(0)
+        assert iobuf.io_buffer.read() == b'\x03\x04'

tests/unit/test_protocol.py

@@ -21,7 +21,7 @@
     PrepareMessage, QueryMessage, ExecuteMessage, UnsupportedOperation,
     _PAGING_OPTIONS_FLAG, _WITH_SERIAL_CONSISTENCY_FLAG,
     _PAGE_SIZE_FLAG, _WITH_PAGING_STATE_FLAG,
-    BatchMessage
+    BatchMessage, BytesReader
 )
 from cassandra.query import BatchType
 from cassandra.marshal import uint32_unpack
@@ -189,3 +189,61 @@ def test_batch_message_with_keyspace(self):
              (b'\x00\x03',),
              (b'\x00\x00\x00\x80',), (b'\x00\x02',), (b'ks',))
         )
+
+
+class BytesReaderTest(unittest.TestCase):
+
+    def test_read_exact_n_bytes(self):
+        reader = BytesReader(b'\x01\x02\x03\x04\x05')
+        assert reader.read(3) == b'\x01\x02\x03'
+
+    def test_read_all_with_negative(self):
+        reader = BytesReader(b'hello')
+        assert reader.read(-1) == b'hello'
+
+    def test_read_all_default(self):
+        reader = BytesReader(b'hello')
+        reader.read(2)
+        assert reader.read(-1) == b'llo'
+
+    def test_sequential_reads(self):
+        reader = BytesReader(b'\x01\x02\x03\x04')
+        assert reader.read(2) == b'\x01\x02'
+        assert reader.read(2) == b'\x03\x04'
+
+    def test_read_past_end_returns_partial(self):
+        reader = BytesReader(b'\x01\x02')
+        assert reader.read(3) == b'\x01\x02'
+
+    def test_read_past_end_after_partial_consume(self):
+        reader = BytesReader(b'\x01\x02\x03')
+        reader.read(2)
+        assert reader.read(2) == b'\x03'
+
+    def test_read_at_end_returns_empty(self):
+        reader = BytesReader(b'\x01\x02')
+        reader.read(2)
+        assert reader.read(1) == b''
+
+    def test_read_zero_bytes(self):
+        reader = BytesReader(b'abc')
+        assert reader.read(0) == b''
+
+    def test_memoryview_input_returns_bytes(self):
+        data = b'\x01\x02\x03\x04'
+        reader = BytesReader(memoryview(data))
+        result = reader.read(2)
+        assert isinstance(result, bytes)
+        assert result == b'\x01\x02'
+
+    def test_memoryview_read_all_returns_bytes(self):
+        data = b'\x01\x02\x03'
+        reader = BytesReader(memoryview(data))
+        result = reader.read(-1)
+        assert isinstance(result, bytes)
+        assert result == b'\x01\x02\x03'
+
+    def test_empty_input(self):
+        reader = BytesReader(b'')
+        assert reader.read(-1) == b''
+        assert reader.read(1) == b''