diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index a7f66d9620..f6e6d7b722 100644 --- a/test/asynchronous/test_client.py +++ b/test/asynchronous/test_client.py @@ -71,7 +71,13 @@ from pymongo.asynchronous.topology import _ErrorContext from pymongo.client_options import ClientOptions from pymongo.common import _UUID_REPRESENTATIONS, CONNECT_TIMEOUT, MIN_SUPPORTED_WIRE_VERSION, has_c -from pymongo.compression_support import _have_snappy, _have_zstd +from pymongo.compression_support import ( + SnappyContext, + ZlibContext, + ZstdContext, + _have_snappy, + _have_zstd, +) from pymongo.driver_info import DriverInfo from pymongo.errors import ( AutoReconnect, @@ -1812,6 +1818,48 @@ def compression_settings(client): # No error await client.pymongo_test.test.find_one() + async def test_compression_commands(self): + # Ensure the compression logic is actually exercised end-to-end by + # sending commands with each available compressor negotiated. + candidates: list[tuple[str, type]] = [("zlib", ZlibContext)] + if _have_snappy(): + candidates.append(("snappy", SnappyContext)) + if _have_zstd(): + candidates.append(("zstd", ZstdContext)) + + for name, ctx_type in candidates: + with self.subTest(compressor=name): + # maxPoolSize=1 ensures the operations below reuse the same + # connection the spy is installed on. + client = await self.async_single_client(compressors=name, maxPoolSize=1) + # Trigger the connection handshake so the compressor is negotiated. + await client.admin.command("ping") + pool = await async_get_pool(client) + async with pool.checkout() as conn: + if conn.compression_context is None: + self.skipTest(f"server did not negotiate {name} compression") + self.assertIsInstance(conn.compression_context, ctx_type) + + # Spy on the compress method to confirm the outgoing message + # is actually compressed. + compressed = [] + original = conn.compression_context.compress + + def spy(data, _original=original, _sink=compressed): + _sink.append(data) + return _original(data) + + conn.compression_context.compress = spy + + # Round-trip a command large enough to compress. The response is + # decompressed via compression_support.decompress. + coll = client.pymongo_test.test_compression + await coll.delete_many({}) + await coll.insert_one({"x": "y" * 1024}) + doc = await coll.find_one({}, {"_id": 0}) + self.assertEqual(doc, {"x": "y" * 1024}) + self.assertTrue(compressed, "compress() was never called") + @async_client_context.require_sync async def test_reset_during_update_pool(self): client = await self.async_rs_or_single_client(minPoolSize=10) diff --git a/test/test_client.py b/test/test_client.py index 8f0da71321..0af6619dd6 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -60,7 +60,13 @@ from pymongo import event_loggers, message, monitoring from pymongo.client_options import ClientOptions from pymongo.common import _UUID_REPRESENTATIONS, CONNECT_TIMEOUT, MIN_SUPPORTED_WIRE_VERSION, has_c -from pymongo.compression_support import _have_snappy, _have_zstd +from pymongo.compression_support import ( + SnappyContext, + ZlibContext, + ZstdContext, + _have_snappy, + _have_zstd, +) from pymongo.driver_info import DriverInfo from pymongo.errors import ( AutoReconnect, @@ -1769,6 +1775,48 @@ def compression_settings(client): # No error client.pymongo_test.test.find_one() + def test_compression_commands(self): + # Ensure the compression logic is actually exercised end-to-end by + # sending commands with each available compressor negotiated. + candidates: list[tuple[str, type]] = [("zlib", ZlibContext)] + if _have_snappy(): + candidates.append(("snappy", SnappyContext)) + if _have_zstd(): + candidates.append(("zstd", ZstdContext)) + + for name, ctx_type in candidates: + with self.subTest(compressor=name): + # maxPoolSize=1 ensures the operations below reuse the same + # connection the spy is installed on. + client = self.single_client(compressors=name, maxPoolSize=1) + # Trigger the connection handshake so the compressor is negotiated. + client.admin.command("ping") + pool = get_pool(client) + with pool.checkout() as conn: + if conn.compression_context is None: + self.skipTest(f"server did not negotiate {name} compression") + self.assertIsInstance(conn.compression_context, ctx_type) + + # Spy on the compress method to confirm the outgoing message + # is actually compressed. + compressed = [] + original = conn.compression_context.compress + + def spy(data, _original=original, _sink=compressed): + _sink.append(data) + return _original(data) + + conn.compression_context.compress = spy + + # Round-trip a command large enough to compress. The response is + # decompressed via compression_support.decompress. + coll = client.pymongo_test.test_compression + coll.delete_many({}) + coll.insert_one({"x": "y" * 1024}) + doc = coll.find_one({}, {"_id": 0}) + self.assertEqual(doc, {"x": "y" * 1024}) + self.assertTrue(compressed, "compress() was never called") + @client_context.require_sync def test_reset_during_update_pool(self): client = self.rs_or_single_client(minPoolSize=10)