From de90b459af7ec53aac36f0eb5aab837c7703877c Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Wed, 20 Nov 2024 16:46:02 -0500 Subject: [PATCH 1/3] Fix GCP + Azure encryption tests --- test/asynchronous/test_encryption.py | 1 + test/test_encryption.py | 1 + 2 files changed, 2 insertions(+) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index a34741c144..8ff1779934 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -1483,6 +1483,7 @@ class AzureGCPEncryptionTestMixin(AsyncEncryptionIntegrationTest): client: AsyncMongoClient async def asyncSetUp(self): + await super().asyncSetUp() self.client = self.simple_client() keyvault = self.client.get_database(self.KEYVAULT_DB).get_collection(self.KEYVAULT_COLL) await create_key_vault(keyvault, self.DEK) diff --git a/test/test_encryption.py b/test/test_encryption.py index 6e6d8ec4d5..b17e193f76 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -1477,6 +1477,7 @@ class AzureGCPEncryptionTestMixin(EncryptionIntegrationTest): client: MongoClient def setUp(self): + super().setUp() self.client = self.simple_client() keyvault = self.client.get_database(self.KEYVAULT_DB).get_collection(self.KEYVAULT_COLL) create_key_vault(keyvault, self.DEK) From 8f6d7385ed1f04f64614de361071c84a5da5b594 Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Tue, 26 Nov 2024 10:41:09 -0500 Subject: [PATCH 2/3] Fix async client unified closing --- test/asynchronous/test_encryption.py | 12 +++++------- test/asynchronous/unified_format.py | 4 ---- test/asynchronous/utils_spec_runner.py | 2 -- test/test_encryption.py | 12 +++++------- test/unified_format.py | 4 ---- test/utils_spec_runner.py | 2 -- 6 files changed, 10 insertions(+), 26 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 8ff1779934..75aec5356d 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -1495,7 +1495,6 @@ async def _test_explicit(self, expectation): async_client_context.client, OPTS, ) - self.addAsyncCleanup(client_encryption.close) ciphertext = await client_encryption.encrypt( "string0", @@ -1521,7 +1520,6 @@ async def _test_automatic(self, expectation_extjson, payload): client = await self.async_rs_or_single_client( auto_encryption_opts=encryption_opts, event_listeners=[insert_listener] ) - self.addAsyncCleanup(client.aclose) coll = client.get_database(encrypted_db).get_collection( encrypted_coll, codec_options=OPTS, write_concern=WriteConcern("majority") @@ -1545,10 +1543,10 @@ async def _test_automatic(self, expectation_extjson, payload): class TestAzureEncryption(AzureGCPEncryptionTestMixin, AsyncEncryptionIntegrationTest): @unittest.skipUnless(any(AZURE_CREDS.values()), "Azure environment credentials are not set") async def asyncSetUp(self): + await super().asyncSetUp() self.KMS_PROVIDER_MAP = {"azure": AZURE_CREDS} self.DEK = json_data(BASE, "custom", "azure-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") - await super().asyncSetUp() async def test_explicit(self): return await self._test_explicit( @@ -1570,10 +1568,10 @@ async def test_automatic(self): class TestGCPEncryption(AzureGCPEncryptionTestMixin, AsyncEncryptionIntegrationTest): @unittest.skipUnless(any(GCP_CREDS.values()), "GCP environment credentials are not set") async def asyncSetUp(self): + await super().asyncSetUp() self.KMS_PROVIDER_MAP = {"gcp": GCP_CREDS} self.DEK = json_data(BASE, "custom", "gcp-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") - await super().asyncSetUp() async def test_explicit(self): return await self._test_explicit( @@ -1595,6 +1593,7 @@ async def test_automatic(self): # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#deadlock-tests class TestDeadlockProse(AsyncEncryptionIntegrationTest): async def asyncSetUp(self): + await super().asyncSetUp() self.client_test = await self.async_rs_or_single_client( maxPoolSize=1, readConcernLevel="majority", w="majority", uuidRepresentation="standard" ) @@ -1627,7 +1626,6 @@ async def asyncSetUp(self): self.ciphertext = await client_encryption.encrypt( "string0", Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_alt_name="local" ) - await client_encryption.close() self.client_listener = OvertCommandListener() self.topology_listener = TopologyEventListener() @@ -1822,6 +1820,7 @@ async def test_case_8(self): # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#14-decryption-events class TestDecryptProse(AsyncEncryptionIntegrationTest): async def asyncSetUp(self): + await super().asyncSetUp() self.client = async_client_context.client await self.client.db.drop_collection("decryption_events") await create_key_vault(self.client.keyvault.datakeys) @@ -2257,6 +2256,7 @@ async def test_06_named_kms_providers_apply_tls_options_kmip(self): # https://github.com/mongodb/specifications/blob/50e26fe/source/client-side-encryption/tests/README.md#unique-index-on-keyaltnames class TestUniqueIndexOnKeyAltNamesProse(AsyncEncryptionIntegrationTest): async def asyncSetUp(self): + await super().asyncSetUp() self.client = async_client_context.client await create_key_vault(self.client.keyvault.datakeys) kms_providers_map = {"local": {"key": LOCAL_MASTER_KEY}} @@ -2606,8 +2606,6 @@ async def AsyncMongoClient(**kwargs): assert isinstance(res["encrypted_indexed"], Binary) assert isinstance(res["encrypted_unindexed"], Binary) - await client_encryption.close() - # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#22-range-explicit-encryption class TestRangeQueryProse(AsyncEncryptionIntegrationTest): diff --git a/test/asynchronous/unified_format.py b/test/asynchronous/unified_format.py index ad02778315..b18b09383e 100644 --- a/test/asynchronous/unified_format.py +++ b/test/asynchronous/unified_format.py @@ -499,10 +499,6 @@ async def asyncSetUp(self): # process file-level runOnRequirements run_on_spec = self.TEST_SPEC.get("runOnRequirements", []) if not await self.should_run_on(run_on_spec): - # Explicitly close async clients here - # to prevent leaky monitor tasks - if not _IS_SYNC: - await async_client_context.client.close() raise unittest.SkipTest(f"{self.__class__.__name__} runOnRequirements not satisfied") # add any special-casing for skipping tests here diff --git a/test/asynchronous/utils_spec_runner.py b/test/asynchronous/utils_spec_runner.py index 9ce9ed6822..b79e5258b5 100644 --- a/test/asynchronous/utils_spec_runner.py +++ b/test/asynchronous/utils_spec_runner.py @@ -692,8 +692,6 @@ async def run_scenario(self, scenario_def, test): self.listener = listener self.pool_listener = pool_listener self.server_listener = server_listener - # Close the client explicitly to avoid having too many threads open. - self.addAsyncCleanup(client.close) # Create session0 and session1. sessions = {} diff --git a/test/test_encryption.py b/test/test_encryption.py index b17e193f76..bebcf69880 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -1489,7 +1489,6 @@ def _test_explicit(self, expectation): client_context.client, OPTS, ) - self.addCleanup(client_encryption.close) ciphertext = client_encryption.encrypt( "string0", @@ -1515,7 +1514,6 @@ def _test_automatic(self, expectation_extjson, payload): client = self.rs_or_single_client( auto_encryption_opts=encryption_opts, event_listeners=[insert_listener] ) - self.addCleanup(client.close) coll = client.get_database(encrypted_db).get_collection( encrypted_coll, codec_options=OPTS, write_concern=WriteConcern("majority") @@ -1539,10 +1537,10 @@ def _test_automatic(self, expectation_extjson, payload): class TestAzureEncryption(AzureGCPEncryptionTestMixin, EncryptionIntegrationTest): @unittest.skipUnless(any(AZURE_CREDS.values()), "Azure environment credentials are not set") def setUp(self): + super().setUp() self.KMS_PROVIDER_MAP = {"azure": AZURE_CREDS} self.DEK = json_data(BASE, "custom", "azure-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") - super().setUp() def test_explicit(self): return self._test_explicit( @@ -1564,10 +1562,10 @@ def test_automatic(self): class TestGCPEncryption(AzureGCPEncryptionTestMixin, EncryptionIntegrationTest): @unittest.skipUnless(any(GCP_CREDS.values()), "GCP environment credentials are not set") def setUp(self): + super().setUp() self.KMS_PROVIDER_MAP = {"gcp": GCP_CREDS} self.DEK = json_data(BASE, "custom", "gcp-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") - super().setUp() def test_explicit(self): return self._test_explicit( @@ -1589,6 +1587,7 @@ def test_automatic(self): # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#deadlock-tests class TestDeadlockProse(EncryptionIntegrationTest): def setUp(self): + super().setUp() self.client_test = self.rs_or_single_client( maxPoolSize=1, readConcernLevel="majority", w="majority", uuidRepresentation="standard" ) @@ -1619,7 +1618,6 @@ def setUp(self): self.ciphertext = client_encryption.encrypt( "string0", Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_alt_name="local" ) - client_encryption.close() self.client_listener = OvertCommandListener() self.topology_listener = TopologyEventListener() @@ -1814,6 +1812,7 @@ def test_case_8(self): # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#14-decryption-events class TestDecryptProse(EncryptionIntegrationTest): def setUp(self): + super().setUp() self.client = client_context.client self.client.db.drop_collection("decryption_events") create_key_vault(self.client.keyvault.datakeys) @@ -2249,6 +2248,7 @@ def test_06_named_kms_providers_apply_tls_options_kmip(self): # https://github.com/mongodb/specifications/blob/50e26fe/source/client-side-encryption/tests/README.md#unique-index-on-keyaltnames class TestUniqueIndexOnKeyAltNamesProse(EncryptionIntegrationTest): def setUp(self): + super().setUp() self.client = client_context.client create_key_vault(self.client.keyvault.datakeys) kms_providers_map = {"local": {"key": LOCAL_MASTER_KEY}} @@ -2590,8 +2590,6 @@ def MongoClient(**kwargs): assert isinstance(res["encrypted_indexed"], Binary) assert isinstance(res["encrypted_unindexed"], Binary) - client_encryption.close() - # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#22-range-explicit-encryption class TestRangeQueryProse(EncryptionIntegrationTest): diff --git a/test/unified_format.py b/test/unified_format.py index a12ef805f2..5cb268a29d 100644 --- a/test/unified_format.py +++ b/test/unified_format.py @@ -498,10 +498,6 @@ def setUp(self): # process file-level runOnRequirements run_on_spec = self.TEST_SPEC.get("runOnRequirements", []) if not self.should_run_on(run_on_spec): - # Explicitly close async clients here - # to prevent leaky monitor tasks - if not _IS_SYNC: - client_context.client.close() raise unittest.SkipTest(f"{self.__class__.__name__} runOnRequirements not satisfied") # add any special-casing for skipping tests here diff --git a/test/utils_spec_runner.py b/test/utils_spec_runner.py index c50d23175b..4508502cd0 100644 --- a/test/utils_spec_runner.py +++ b/test/utils_spec_runner.py @@ -689,8 +689,6 @@ def run_scenario(self, scenario_def, test): self.listener = listener self.pool_listener = pool_listener self.server_listener = server_listener - # Close the client explicitly to avoid having too many threads open. - self.addCleanup(client.close) # Create session0 and session1. sessions = {} From 619a2ce5fba1881d5726826c984d1fd650de7372 Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Tue, 26 Nov 2024 11:19:06 -0500 Subject: [PATCH 3/3] Fix batch bulk tests --- .evergreen/config.yml | 2 +- test/asynchronous/test_encryption.py | 18 +++++++++++------- test/test_encryption.py | 18 +++++++++++------- 3 files changed, 23 insertions(+), 15 deletions(-) diff --git a/.evergreen/config.yml b/.evergreen/config.yml index 59b8a543fd..7ca3a72b1a 100644 --- a/.evergreen/config.yml +++ b/.evergreen/config.yml @@ -281,7 +281,7 @@ functions: "run tests": - command: subprocess.exec params: - include_expansions_in_env: ["TEST_DATA_LAKE", "AUTH", "SSL", "TEST_INDEX_MANAGEMENT", "CRYPT_SHARED_LIB_PATH", "test_encryption", "test_encryption_pyopenssl", "test_crypt_shared", "test_pyopenssl", "test_loadbalancer", "test_serverless", "ORCHESTRATION_FILE"] + include_expansions_in_env: ["TEST_DATA_LAKE", "PYTHON_BINARY", "AUTH", "SSL", "TEST_INDEX_MANAGEMENT", "CRYPT_SHARED_LIB_PATH", "test_encryption", "test_encryption_pyopenssl", "test_crypt_shared", "test_pyopenssl", "test_loadbalancer", "test_serverless", "ORCHESTRATION_FILE"] binary: bash working_dir: "src" args: diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 75aec5356d..048db2d501 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -1234,7 +1234,9 @@ async def test_03_bulk_batch_split(self): doc2 = {"_id": "over_2mib_2", "unencrypted": "a" * _2_MiB} self.listener.reset() await self.coll_encrypted.bulk_write([InsertOne(doc1), InsertOne(doc2)]) - self.assertEqual(self.listener.started_command_names(), ["insert", "insert"]) + self.assertEqual( + len([c for c in self.listener.started_command_names() if c == "insert"]), 2 + ) async def test_04_bulk_batch_split(self): limits_doc = json_data("limits", "limits-doc.json") @@ -1244,7 +1246,9 @@ async def test_04_bulk_batch_split(self): doc2.update(limits_doc) self.listener.reset() await self.coll_encrypted.bulk_write([InsertOne(doc1), InsertOne(doc2)]) - self.assertEqual(self.listener.started_command_names(), ["insert", "insert"]) + self.assertEqual( + len([c for c in self.listener.started_command_names() if c == "insert"]), 2 + ) async def test_05_insert_succeeds_just_under_16MiB(self): doc = {"_id": "under_16mib", "unencrypted": "a" * (_16_MiB - 2000)} @@ -1482,13 +1486,12 @@ class AzureGCPEncryptionTestMixin(AsyncEncryptionIntegrationTest): KEYVAULT_COLL = "datakeys" client: AsyncMongoClient - async def asyncSetUp(self): - await super().asyncSetUp() - self.client = self.simple_client() + async def _setup(self): keyvault = self.client.get_database(self.KEYVAULT_DB).get_collection(self.KEYVAULT_COLL) await create_key_vault(keyvault, self.DEK) async def _test_explicit(self, expectation): + await self._setup() client_encryption = self.create_client_encryption( self.KMS_PROVIDER_MAP, # type: ignore[arg-type] ".".join([self.KEYVAULT_DB, self.KEYVAULT_COLL]), @@ -1506,6 +1509,7 @@ async def _test_explicit(self, expectation): self.assertEqual(await client_encryption.decrypt(ciphertext), "string0") async def _test_automatic(self, expectation_extjson, payload): + await self._setup() encrypted_db = "db" encrypted_coll = "coll" keyvault_namespace = ".".join([self.KEYVAULT_DB, self.KEYVAULT_COLL]) @@ -1543,10 +1547,10 @@ async def _test_automatic(self, expectation_extjson, payload): class TestAzureEncryption(AzureGCPEncryptionTestMixin, AsyncEncryptionIntegrationTest): @unittest.skipUnless(any(AZURE_CREDS.values()), "Azure environment credentials are not set") async def asyncSetUp(self): - await super().asyncSetUp() self.KMS_PROVIDER_MAP = {"azure": AZURE_CREDS} self.DEK = json_data(BASE, "custom", "azure-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") + await super().asyncSetUp() async def test_explicit(self): return await self._test_explicit( @@ -1568,10 +1572,10 @@ async def test_automatic(self): class TestGCPEncryption(AzureGCPEncryptionTestMixin, AsyncEncryptionIntegrationTest): @unittest.skipUnless(any(GCP_CREDS.values()), "GCP environment credentials are not set") async def asyncSetUp(self): - await super().asyncSetUp() self.KMS_PROVIDER_MAP = {"gcp": GCP_CREDS} self.DEK = json_data(BASE, "custom", "gcp-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") + await super().asyncSetUp() async def test_explicit(self): return await self._test_explicit( diff --git a/test/test_encryption.py b/test/test_encryption.py index bebcf69880..cb8bcb74d6 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -1230,7 +1230,9 @@ def test_03_bulk_batch_split(self): doc2 = {"_id": "over_2mib_2", "unencrypted": "a" * _2_MiB} self.listener.reset() self.coll_encrypted.bulk_write([InsertOne(doc1), InsertOne(doc2)]) - self.assertEqual(self.listener.started_command_names(), ["insert", "insert"]) + self.assertEqual( + len([c for c in self.listener.started_command_names() if c == "insert"]), 2 + ) def test_04_bulk_batch_split(self): limits_doc = json_data("limits", "limits-doc.json") @@ -1240,7 +1242,9 @@ def test_04_bulk_batch_split(self): doc2.update(limits_doc) self.listener.reset() self.coll_encrypted.bulk_write([InsertOne(doc1), InsertOne(doc2)]) - self.assertEqual(self.listener.started_command_names(), ["insert", "insert"]) + self.assertEqual( + len([c for c in self.listener.started_command_names() if c == "insert"]), 2 + ) def test_05_insert_succeeds_just_under_16MiB(self): doc = {"_id": "under_16mib", "unencrypted": "a" * (_16_MiB - 2000)} @@ -1476,13 +1480,12 @@ class AzureGCPEncryptionTestMixin(EncryptionIntegrationTest): KEYVAULT_COLL = "datakeys" client: MongoClient - def setUp(self): - super().setUp() - self.client = self.simple_client() + def _setup(self): keyvault = self.client.get_database(self.KEYVAULT_DB).get_collection(self.KEYVAULT_COLL) create_key_vault(keyvault, self.DEK) def _test_explicit(self, expectation): + self._setup() client_encryption = self.create_client_encryption( self.KMS_PROVIDER_MAP, # type: ignore[arg-type] ".".join([self.KEYVAULT_DB, self.KEYVAULT_COLL]), @@ -1500,6 +1503,7 @@ def _test_explicit(self, expectation): self.assertEqual(client_encryption.decrypt(ciphertext), "string0") def _test_automatic(self, expectation_extjson, payload): + self._setup() encrypted_db = "db" encrypted_coll = "coll" keyvault_namespace = ".".join([self.KEYVAULT_DB, self.KEYVAULT_COLL]) @@ -1537,10 +1541,10 @@ def _test_automatic(self, expectation_extjson, payload): class TestAzureEncryption(AzureGCPEncryptionTestMixin, EncryptionIntegrationTest): @unittest.skipUnless(any(AZURE_CREDS.values()), "Azure environment credentials are not set") def setUp(self): - super().setUp() self.KMS_PROVIDER_MAP = {"azure": AZURE_CREDS} self.DEK = json_data(BASE, "custom", "azure-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") + super().setUp() def test_explicit(self): return self._test_explicit( @@ -1562,10 +1566,10 @@ def test_automatic(self): class TestGCPEncryption(AzureGCPEncryptionTestMixin, EncryptionIntegrationTest): @unittest.skipUnless(any(GCP_CREDS.values()), "GCP environment credentials are not set") def setUp(self): - super().setUp() self.KMS_PROVIDER_MAP = {"gcp": GCP_CREDS} self.DEK = json_data(BASE, "custom", "gcp-dek.json") self.SCHEMA_MAP = json_data(BASE, "custom", "azure-gcp-schema.json") + super().setUp() def test_explicit(self): return self._test_explicit(