PYTHON-5143 Support auto encryption in unified tests

Author: blink1073

.evergreen/remove-unimplemented-tests.sh

@@ -3,11 +3,6 @@ PYMONGO=$(dirname "$(cd "$(dirname "$0")" || exit; pwd)")
 
 rm $PYMONGO/test/transactions/legacy/errors-client.json  # PYTHON-1894
 rm $PYMONGO/test/connection_monitoring/wait-queue-fairness.json  # PYTHON-1873
-rm $PYMONGO/test/client-side-encryption/spec/unified/fle2v2-BypassQueryAnalysis.json  # PYTHON-5143
-rm $PYMONGO/test/client-side-encryption/spec/unified/fle2v2-EncryptedFields-vs-EncryptedFieldsMap.json  # PYTHON-5143
-rm $PYMONGO/test/client-side-encryption/spec/unified/localSchema.json  # PYTHON-5143
-rm $PYMONGO/test/client-side-encryption/spec/unified/maxWireVersion.json  # PYTHON-5143
-rm $PYMONGO/test/unified-test-format/valid-pass/poc-queryable-encryption.json  # PYTHON-5143
 rm $PYMONGO/test/discovery_and_monitoring/unified/pool-clear-application-error.json  # PYTHON-4918
 rm $PYMONGO/test/discovery_and_monitoring/unified/pool-clear-checkout-error.json  # PYTHON-4918
 rm $PYMONGO/test/discovery_and_monitoring/unified/pool-clear-min-pool-size-error.json  # PYTHON-4918

test/asynchronous/helpers.py

@@ -116,6 +116,15 @@
 }
 KMIP_CREDS = {"endpoint": os.environ.get("FLE_KMIP_ENDPOINT", "localhost:5698")}
 
+ALL_KMS_PROVIDERS = dict(
+    aws=AWS_CREDS,
+    azure=AZURE_CREDS,
+    gcp=GCP_CREDS,
+    local=dict(key=LOCAL_MASTER_KEY),
+    kmip=KMIP_CREDS,
+)
+DEFAULT_KMS_TLS = dict(kmip=dict(tlsCAFile=CA_PEM, tlsCertificateKeyFile=CLIENT_PEM))
+
 # Ensure Evergreen metadata doesn't result in truncation
 os.environ.setdefault("MONGOB_LOG_MAX_DOCUMENT_LENGTH", "2000")
 

test/asynchronous/test_encryption.py

@@ -54,18 +54,20 @@
 from test import (
     unittest,
 )
-from test.asynchronous.test_bulk import AsyncBulkTestBase
-from test.asynchronous.unified_format import generate_test_classes
-from test.asynchronous.utils_spec_runner import AsyncSpecRunner
-from test.helpers import (
+from test.asynchronous.helpers import (
+    ALL_KMS_PROVIDERS,
     AWS_CREDS,
     AZURE_CREDS,
     CA_PEM,
     CLIENT_PEM,
+    DEFAULT_KMS_TLS,
     GCP_CREDS,
     KMIP_CREDS,
     LOCAL_MASTER_KEY,
 )
+from test.asynchronous.test_bulk import AsyncBulkTestBase
+from test.asynchronous.unified_format import generate_test_classes
+from test.asynchronous.utils_spec_runner import AsyncSpecRunner
 from test.utils_shared import (
     AllowListEventListener,
     OvertCommandListener,
@@ -204,7 +206,7 @@ async def test_init_kms_tls_options(self):
         opts = AutoEncryptionOpts(
             {},
             "k.d",
-            kms_tls_options={"kmip": {"tlsCAFile": CA_PEM, "tlsCertificateKeyFile": CLIENT_PEM}},
+            kms_tls_options=DEFAULT_KMS_TLS,
         )
         _kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
         ctx = _kms_ssl_contexts["kmip"]
@@ -626,7 +628,6 @@ async def test_with_statement(self):
     "accessKeyId": os.environ.get("CSFLE_AWS_TEMP_ACCESS_KEY_ID", ""),
     "secretAccessKey": os.environ.get("CSFLE_AWS_TEMP_SECRET_ACCESS_KEY", ""),
 }
-KMS_TLS_OPTS = {"kmip": {"tlsCAFile": CA_PEM, "tlsCertificateKeyFile": CLIENT_PEM}}
 
 
 class AsyncTestSpec(AsyncSpecRunner):
@@ -663,7 +664,7 @@ def parse_auto_encrypt_opts(self, opts):
                 self.skipTest("GCP environment credentials are not set")
         if "kmip" in kms_providers:
             kms_providers["kmip"] = KMIP_CREDS
-            opts["kms_tls_options"] = KMS_TLS_OPTS
+            opts["kms_tls_options"] = DEFAULT_KMS_TLS
         if "key_vault_namespace" not in opts:
             opts["key_vault_namespace"] = "keyvault.datakeys"
         if "extra_options" in opts:
@@ -757,14 +758,6 @@ async def run_scenario(self):
     )
 
 # Prose Tests
-ALL_KMS_PROVIDERS = {
-    "aws": AWS_CREDS,
-    "azure": AZURE_CREDS,
-    "gcp": GCP_CREDS,
-    "kmip": KMIP_CREDS,
-    "local": {"key": LOCAL_MASTER_KEY},
-}
-
 LOCAL_KEY_ID = Binary(base64.b64decode(b"LOCALAAAAAAAAAAAAAAAAA=="), UUID_SUBTYPE)
 AWS_KEY_ID = Binary(base64.b64decode(b"AWSAAAAAAAAAAAAAAAAAAA=="), UUID_SUBTYPE)
 AZURE_KEY_ID = Binary(base64.b64decode(b"AZUREAAAAAAAAAAAAAAAAA=="), UUID_SUBTYPE)
@@ -851,13 +844,17 @@ async def asyncSetUp(self):
             self.KMS_PROVIDERS,
             "keyvault.datakeys",
             schema_map=schemas,
-            kms_tls_options=KMS_TLS_OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
         )
         self.client_encrypted = await self.async_rs_or_single_client(
             auto_encryption_opts=opts, uuidRepresentation="standard"
         )
         self.client_encryption = self.create_client_encryption(
-            self.KMS_PROVIDERS, "keyvault.datakeys", self.client, OPTS, kms_tls_options=KMS_TLS_OPTS
+            self.KMS_PROVIDERS,
+            "keyvault.datakeys",
+            self.client,
+            OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
         )
         self.listener.reset()
 
@@ -1066,7 +1063,7 @@ async def _test_corpus(self, opts):
             "keyvault.datakeys",
             async_client_context.client,
             OPTS,
-            kms_tls_options=KMS_TLS_OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
         )
 
         corpus = self.fix_up_curpus(json_data("corpus", "corpus.json"))
@@ -1158,7 +1155,7 @@ async def _test_corpus(self, opts):
 
     async def test_corpus(self):
         opts = AutoEncryptionOpts(
-            self.kms_providers(), "keyvault.datakeys", kms_tls_options=KMS_TLS_OPTS
+            self.kms_providers(), "keyvault.datakeys", kms_tls_options=DEFAULT_KMS_TLS
         )
         await self._test_corpus(opts)
 
@@ -1169,7 +1166,7 @@ async def test_corpus_local_schema(self):
             self.kms_providers(),
             "keyvault.datakeys",
             schema_map=schemas,
-            kms_tls_options=KMS_TLS_OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
         )
         await self._test_corpus(opts)
 
@@ -1300,7 +1297,7 @@ async def asyncSetUp(self):
             key_vault_namespace="keyvault.datakeys",
             key_vault_client=async_client_context.client,
             codec_options=OPTS,
-            kms_tls_options=KMS_TLS_OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
         )
 
         kms_providers_invalid = copy.deepcopy(kms_providers)
@@ -1312,7 +1309,7 @@ async def asyncSetUp(self):
             key_vault_namespace="keyvault.datakeys",
             key_vault_client=async_client_context.client,
             codec_options=OPTS,
-            kms_tls_options=KMS_TLS_OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
         )
         self._kmip_host_error = None
         self._invalid_host_error = None
@@ -2752,7 +2749,7 @@ async def run_test(self, src_provider, dst_provider):
             key_vault_client=self.client,
             key_vault_namespace="keyvault.datakeys",
             kms_providers=ALL_KMS_PROVIDERS,
-            kms_tls_options=KMS_TLS_OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
             codec_options=OPTS,
         )
 
@@ -2772,7 +2769,7 @@ async def run_test(self, src_provider, dst_provider):
             key_vault_client=client2,
             key_vault_namespace="keyvault.datakeys",
             kms_providers=ALL_KMS_PROVIDERS,
-            kms_tls_options=KMS_TLS_OPTS,
+            kms_tls_options=DEFAULT_KMS_TLS,
             codec_options=OPTS,
         )
 

test/asynchronous/unified_format.py

@@ -37,6 +37,7 @@
 )
 from test.asynchronous.utils import async_get_pool, flaky
 from test.asynchronous.utils_spec_runner import SpecRunnerTask
+from test.helpers import ALL_KMS_PROVIDERS, DEFAULT_KMS_TLS
 from test.unified_format_shared import (
     KMS_TLS_OPTS,
     PLACEHOLDER_MAP,
@@ -61,6 +62,8 @@
 from test.version import Version
 from typing import Any, Dict, List, Mapping, Optional
 
+import pytest
+
 import pymongo
 from bson import SON, json_util
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
@@ -76,7 +79,7 @@
 from pymongo.asynchronous.encryption import AsyncClientEncryption
 from pymongo.asynchronous.helpers import anext
 from pymongo.driver_info import DriverInfo
-from pymongo.encryption_options import _HAVE_PYMONGOCRYPT
+from pymongo.encryption_options import _HAVE_PYMONGOCRYPT, AutoEncryptionOpts
 from pymongo.errors import (
     AutoReconnect,
     BulkWriteError,
@@ -259,6 +262,23 @@ async def _create_entity(self, entity_spec, uri=None):
             kwargs: dict = {}
             observe_events = spec.get("observeEvents", [])
 
+            if "autoEncryptOpts" in spec:
+                auto_encrypt_opts = spec["autoEncryptOpts"]
+                auto_encrypt_kwargs: dict = dict(kms_tls_options=DEFAULT_KMS_TLS)
+                kms_providers = ALL_KMS_PROVIDERS.copy()
+                key_vault_namespace = auto_encrypt_opts.pop("keyVaultNamespace")
+                for provider_name, provider_value in auto_encrypt_opts.pop("kmsProviders").items():
+                    kms_providers[provider_name].update(provider_value)
+                extra_opts = auto_encrypt_opts.pop("extraOptions", {})
+                for key, value in extra_opts.items():
+                    auto_encrypt_kwargs[camel_to_snake(key)] = value
+                for key, value in auto_encrypt_opts.items():
+                    auto_encrypt_kwargs[camel_to_snake(key)] = value
+                auto_encryption_opts = AutoEncryptionOpts(
+                    kms_providers, key_vault_namespace, **auto_encrypt_kwargs
+                )
+                kwargs["auto_encryption_opts"] = auto_encryption_opts
+
             # The unified tests use topologyOpeningEvent, we use topologyOpenedEvent
             for i in range(len(observe_events)):
                 if "topologyOpeningEvent" == observe_events[i]:
@@ -430,7 +450,7 @@ class UnifiedSpecTestMixinV1(AsyncIntegrationTest):
     a class attribute ``TEST_SPEC``.
     """
 
-    SCHEMA_VERSION = Version.from_string("1.22")
+    SCHEMA_VERSION = Version.from_string("1.23")
     RUN_ON_LOAD_BALANCER = True
     TEST_SPEC: Any
     TEST_PATH = ""  # This gets filled in by generate_test_classes
@@ -1516,7 +1536,14 @@ class SpecTestBase(with_metaclass(UnifiedSpecTestMeta)):  # type: ignore
             TEST_SPEC = test_spec
             EXPECTED_FAILURES = expected_failures
 
-        return SpecTestBase
+        base = SpecTestBase
+
+        # Add "encryption" marker if the "csfle" runOnRequirement is set.
+        for req in test_spec.get("runOnRequirements", []):
+            if req.get("csfle", False):
+                base = pytest.mark.encryption(base)
+
+        return base
 
     for dirpath, _, filenames in os.walk(test_path):
         dirname = os.path.split(dirpath)[-1]