Revert "Noah feedback"

This reverts commit 4e02b3f.

Author: aclark4life

test/asynchronous/test_async_network_layer.py

@@ -30,9 +30,7 @@
 from pymongo.network_layer import PyMongoProtocol, _async_socket_receive
 
 
-def _make_protocol(timeout=None):
-    # PyMongoProtocol.__init__ calls asyncio.get_running_loop(), so this helper
-    # must be called from inside an async test method.
+async def _make_protocol(timeout=None):
     protocol = PyMongoProtocol(timeout=timeout)
     mock_transport = MagicMock()
     mock_transport.is_closing.return_value = False
@@ -45,88 +43,109 @@ def _make_header(length, request_id, response_to, op_code):
 
 
 class TestPyMongoProtocol(AsyncUnitTest):
-    def _make_proto_with_header(self, header_bytes, max_size=MAX_MESSAGE_SIZE):
-        protocol = _make_protocol()
+    async def _make_proto_with_header(self, header_bytes, max_size=MAX_MESSAGE_SIZE):
+        protocol = await _make_protocol()
         protocol._max_message_size = max_size
         protocol._header = memoryview(bytearray(header_bytes))
         return protocol
 
-    def test_normal_op_msg(self):
+    async def test_normal_op_msg(self):
         header = _make_header(length=32, request_id=1, response_to=99, op_code=2013)
-        protocol = self._make_proto_with_header(header)
+        protocol = await self._make_proto_with_header(header)
         body_len, op_code, response_to, expecting_compression = protocol.process_header()
         self.assertEqual(body_len, 16)
         self.assertEqual(op_code, 2013)
         self.assertEqual(response_to, 99)
         self.assertFalse(expecting_compression)
 
-    def test_op_compressed(self):
+    async def test_op_compressed(self):
         # OP_COMPRESSED=2012; process_header strips the 9-byte compression sub-header
         # (op code + uncompressed size + compressor id), then the 16-byte standard header.
         # length=35 → after compression sub-header: 26 → body: 10
         header = _make_header(length=35, request_id=1, response_to=0, op_code=2012)
-        protocol = self._make_proto_with_header(header)
+        protocol = await self._make_proto_with_header(header)
         body_len, op_code, _response_to, expecting_compression = protocol.process_header()
         self.assertEqual(body_len, 10)
         self.assertEqual(op_code, 2012)
         self.assertTrue(expecting_compression)
 
-    def test_op_compressed_length_too_small_raises(self):
+    async def test_op_compressed_length_too_small_raises(self):
         header = _make_header(length=25, request_id=1, response_to=0, op_code=2012)
-        protocol = self._make_proto_with_header(header)
+        protocol = await self._make_proto_with_header(header)
         with self.assertRaises(ProtocolError):
             protocol.process_header()
 
-    def test_non_compressed_length_too_small_raises(self):
+    async def test_non_compressed_length_too_small_raises(self):
         header = _make_header(length=16, request_id=1, response_to=0, op_code=2013)
-        protocol = self._make_proto_with_header(header)
+        protocol = await self._make_proto_with_header(header)
         with self.assertRaises(ProtocolError):
             protocol.process_header()
 
-    def test_length_exceeds_max_raises(self):
+    async def test_length_exceeds_max_raises(self):
         header = _make_header(
             length=MAX_MESSAGE_SIZE + 1, request_id=1, response_to=0, op_code=2013
         )
-        protocol = self._make_proto_with_header(header)
+        protocol = await self._make_proto_with_header(header)
         with self.assertRaises(ProtocolError):
             protocol.process_header()
 
-    def test_op_reply_op_code(self):
+    async def test_op_reply_op_code(self):
         header = _make_header(length=20, request_id=0, response_to=0, op_code=1)
-        protocol = self._make_proto_with_header(header)
+        protocol = await self._make_proto_with_header(header)
         body_len, op_code, _response_to, expecting_compression = protocol.process_header()
         self.assertEqual(body_len, 4)
         self.assertEqual(op_code, 1)
         self.assertFalse(expecting_compression)
 
-    def test_compression_header_snappy_compressor_id(self):
-        protocol = _make_protocol()
+    async def test_compression_header_snappy_compressor_id(self):
+        protocol = await _make_protocol()
         # <iiB: little-endian, i32 op code=2013, i32 uncompressed size=0, u8 compressor id=1 (snappy)
         data = struct.pack("<iiB", 2013, 0, 1)
         protocol._compression_header = memoryview(bytearray(data))
         op_code, compressor_id = protocol.process_compression_header()
         self.assertEqual(op_code, 2013)
         self.assertEqual(compressor_id, 1)
 
-    def test_compression_header_zlib_compressor_id(self):
-        protocol = _make_protocol()
+    async def test_compression_header_zlib_compressor_id(self):
+        protocol = await _make_protocol()
         data = struct.pack("<iiB", 2013, 0, 2)
         protocol._compression_header = memoryview(bytearray(data))
         _, compressor_id = protocol.process_compression_header()
         self.assertEqual(compressor_id, 2)
 
-    def test_close_aborts_transport(self):
-        protocol = _make_protocol()
+    async def test_message_complete_resolves_pending_future(self):
+        protocol = await _make_protocol()
+        protocol._expecting_header = False
+        protocol._expecting_compression = False
+        protocol._message_size = 10
+        protocol._message = memoryview(bytearray(10))
+        protocol._message_index = 0
+        protocol._op_code = 2013
+        protocol._compressor_id = None
+        protocol._response_to = 42
+
+        future = asyncio.get_running_loop().create_future()
+        protocol._pending_messages.append(future)
+
+        protocol.buffer_updated(10)
+        self.assertTrue(future.done())
+        op_code, compressor_id, response_to, _ = future.result()
+        self.assertEqual(op_code, 2013)
+        self.assertIsNone(compressor_id)
+        self.assertEqual(response_to, 42)
+
+    async def test_close_aborts_transport(self):
+        protocol = await _make_protocol()
         protocol.close()
         self.assertTrue(protocol.transport.abort.called)
 
-    def test_connection_lost_twice_does_not_raise(self):
-        protocol = _make_protocol()
+    async def test_connection_lost_twice_does_not_raise(self):
+        protocol = await _make_protocol()
         protocol.connection_lost(None)
         protocol.connection_lost(None)
 
     async def test_close_with_exception_propagates_to_pending(self):
-        protocol = _make_protocol()
+        protocol = await _make_protocol()
         future = asyncio.get_running_loop().create_future()
         protocol._pending_messages.append(future)
         exc = OSError("connection reset")
@@ -136,6 +155,31 @@ async def test_close_with_exception_propagates_to_pending(self):
 
 
 class TestAsyncSocketReceive(AsyncUnitTest):
+    async def test_reads_data_in_multiple_chunks(self):
+        # Covers the loop in _async_socket_receive that accumulates short reads
+        # until the requested length has been received.
+        data = b"abcdefgh"
+        length = len(data)
+        chunk1, chunk2 = data[:4], data[4:]
+        mock_socket = MagicMock()
+        loop = asyncio.get_running_loop()
+        calls = 0
+
+        async def fake_recv_into(sock, buf):
+            nonlocal calls
+            if calls == 0:
+                buf[: len(chunk1)] = chunk1
+                calls += 1
+                return len(chunk1)
+            buf[: len(chunk2)] = chunk2
+            calls += 1
+            return len(chunk2)
+
+        with patch.object(loop, "sock_recv_into", new=AsyncMock(side_effect=fake_recv_into)):
+            result = await _async_socket_receive(mock_socket, length, loop)
+        self.assertEqual(bytes(result), data)
+        self.assertEqual(calls, 2)
+
     async def test_raises_on_connection_closed(self):
         # Covers the explicit `raise OSError("connection closed")` branch when
         # sock_recv_into returns 0.

test/test_network_layer.py

@@ -117,6 +117,29 @@ def test_unknown_opcode_raises(self):
 
 
 class TestReceiveData(UnitTest):
+    def test_reads_data_in_multiple_chunks(self):
+        # Covers the loop in receive_data that accumulates short reads until the
+        # requested length has been received.
+        data = b"abcdefgh"
+        chunk1, chunk2 = data[:4], data[4:]
+        conn = _make_conn()
+        calls = 0
+
+        def fake_recv_into(buf):
+            nonlocal calls
+            if calls == 0:
+                buf[: len(chunk1)] = chunk1
+                calls += 1
+                return len(chunk1)
+            buf[: len(chunk2)] = chunk2
+            calls += 1
+            return len(chunk2)
+
+        conn.conn.recv_into.side_effect = fake_recv_into
+        result = network_layer.receive_data(conn, len(data), deadline=None)
+        self.assertEqual(bytes(result), data)
+        self.assertEqual(calls, 2)
+
     def test_raises_on_connection_closed(self):
         # Covers the explicit `raise OSError("connection closed")` branch when
         # recv_into returns 0.