PYTHON-2267 Test passing UUID to encrypt_expression (#1510)

Author: ShaneHarvey

doc/changelog.rst

@@ -18,6 +18,9 @@ PyMongo 4.7 brings a number of improvements including:
   Named KMS providers enables more than one of each KMS provider type to be configured.
   See the docstring for :class:`~pymongo.encryption_options.AutoEncryptionOpts`.
   Note that named KMS providers requires pymongocrypt >=1.9 and libmongocrypt >=1.9.
+- :meth:`~pymongo.encryption.ClientEncryption.encrypt` and
+  :meth:`~pymongo.encryption.ClientEncryption.encrypt_expression` now allow ``key_id``
+  to be passed in as a :class:`uuid.UUID`.
 - Fixed a bug where :class:`~bson.int64.Int64` instances could not always be encoded by `orjson`_. The following now
   works::
 

pymongo/encryption.py

@@ -831,7 +831,7 @@ def encrypt(
         :return: The encrypted value, a :class:`~bson.binary.Binary` with subtype 6.
 
         .. versionchanged:: 4.7
-            ``key_id`` can now be passed in as a :class:`uuid.UUID`.
+           ``key_id`` can now be passed in as a :class:`uuid.UUID`.
 
         .. versionchanged:: 4.2
            Added the `query_type` and `contention_factor` parameters.
@@ -883,7 +883,7 @@ def encrypt_expression(
         :return: The encrypted expression, a :class:`~bson.RawBSONDocument`.
 
         .. versionchanged:: 4.7
-            ``key_id`` can now be passed in as a :class:`uuid.UUID`.
+           ``key_id`` can now be passed in as a :class:`uuid.UUID`.
 
         .. versionadded:: 4.4
         """

test/test_encryption.py

@@ -2623,10 +2623,12 @@ def setUp(self):
         self.db = self.encrypted_client.db
         self.addCleanup(self.encrypted_client.close)
 
-    def run_expression_find(self, name, expression, expected_elems, range_opts, use_expr=False):
+    def run_expression_find(
+        self, name, expression, expected_elems, range_opts, use_expr=False, key_id=None
+    ):
         find_payload = self.client_encryption.encrypt_expression(
             expression=expression,
-            key_id=self.key1_id,
+            key_id=key_id or self.key1_id,
             algorithm=Algorithm.RANGEPREVIEW,
             query_type=QueryType.RANGEPREVIEW,
             contention_factor=0,
@@ -2668,16 +2670,20 @@ def encrypt_and_cast(i):
         self.assertEqual(self.client_encryption.decrypt(insert_payload), cast_func(6))
 
         # Case 2.
+        expression = {
+            "$and": [
+                {f"encrypted{name}": {"$gte": cast_func(6)}},
+                {f"encrypted{name}": {"$lte": cast_func(200)}},
+            ]
+        }
+        self.run_expression_find(name, expression, [cast_func(i) for i in [6, 30, 200]], range_opts)
+        # Case 2, with UUID key_id
         self.run_expression_find(
             name,
-            {
-                "$and": [
-                    {f"encrypted{name}": {"$gte": cast_func(6)}},
-                    {f"encrypted{name}": {"$lte": cast_func(200)}},
-                ]
-            },
+            expression,
             [cast_func(i) for i in [6, 30, 200]],
             range_opts,
+            key_id=self.key1_id.as_uuid(),
         )
 
         # Case 3.