PYTHON-1952 Test ClientEncryption.close

Support using ClientEncryption in a with-statement.

Author: ShaneHarvey

pymongo/encryption.py

@@ -44,6 +44,7 @@
 
 from pymongo.errors import (ConfigurationError,
                             EncryptionError,
+                            InvalidOperation,
                             ServerSelectionTimeoutError)
 from pymongo.mongo_client import MongoClient
 from pymongo.pool import _configured_socket, PoolOptions
@@ -387,7 +388,6 @@ def __init__(self, kms_providers, key_vault_namespace, key_vault_client,
         self._encryption = ExplicitEncrypter(
             self._io_callbacks, MongoCryptOptions(kms_providers, None))
 
-    @_wrap_encryption_errors
     def create_data_key(self, kms_provider, master_key=None,
                         key_alt_names=None):
         """Create and insert a new data key into the key vault collection.
@@ -419,8 +419,11 @@ def create_data_key(self, kms_provider, master_key=None,
         :Returns:
           The ``_id`` of the created data key document.
         """
-        return self._encryption.create_data_key(
-            kms_provider, master_key=master_key, key_alt_names=key_alt_names)
+        self._check_closed()
+        with _wrap_encryption_errors_ctx():
+            return self._encryption.create_data_key(
+                kms_provider, master_key=master_key,
+                key_alt_names=key_alt_names)
 
     def encrypt(self, value, algorithm, key_id=None, key_alt_name=None):
         """Encrypt a BSON value with a given key and algorithm.
@@ -440,6 +443,7 @@ def encrypt(self, value, algorithm, key_id=None, key_alt_name=None):
         :Returns:
           The encrypted value, a :class:`~bson.binary.Binary` with subtype 6.
         """
+        self._check_closed()
         if (key_id is not None and not (
                 isinstance(key_id, Binary) and
                 key_id.subtype == UUID_SUBTYPE)):
@@ -462,6 +466,7 @@ def decrypt(self, value):
         :Returns:
           The decrypted BSON value.
         """
+        self._check_closed()
         if not (isinstance(value, Binary) and value.subtype == 6):
             raise TypeError(
                 'value to decrypt must be a bson.binary.Binary with subtype 6')
@@ -472,9 +477,29 @@ def decrypt(self, value):
             return decode(decrypted_doc,
                           codec_options=self._codec_options)['v']
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def _check_closed(self):
+        if self._encryption is None:
+            raise InvalidOperation("Cannot use closed ClientEncryption")
+
     def close(self):
-        """Release resources."""
-        self._io_callbacks.close()
-        self._encryption.close()
-        self._io_callbacks = None
-        self._encryption = None
+        """Release resources.
+
+        Note that using this class in a with-statement will automatically call
+        :meth:`close`::
+
+            with ClientEncryption(...) as client_encryption:
+                encrypted = client_encryption.encrypt(value, ...)
+                decrypted = client_encryption.decrypt(encrypted)
+
+        """
+        if self._io_callbacks:
+            self._io_callbacks.close()
+            self._encryption.close()
+            self._io_callbacks = None
+            self._encryption = None

test/test_encryption.py

@@ -37,6 +37,7 @@
 
 from pymongo.errors import (ConfigurationError,
                             EncryptionError,
+                            InvalidOperation,
                             OperationFailure)
 from pymongo.encryption import (Algorithm,
                                 ClientEncryption)
@@ -366,6 +367,31 @@ def test_codec_options(self):
         self.assertNotEqual(
             client_encryption.decrypt(encrypted_legacy), value)
 
+    def test_close(self):
+        client_encryption = ClientEncryption(
+            KMS_PROVIDERS, 'admin.datakeys', client_context.client, OPTS)
+        client_encryption.close()
+        # Close can be called multiple times.
+        client_encryption.close()
+        algo = Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic
+        msg = 'Cannot use closed ClientEncryption'
+        with self.assertRaisesRegex(InvalidOperation, msg):
+            client_encryption.create_data_key('local')
+        with self.assertRaisesRegex(InvalidOperation, msg):
+            client_encryption.encrypt('val', algo, key_alt_name='name')
+        with self.assertRaisesRegex(InvalidOperation, msg):
+            client_encryption.decrypt(Binary(b'', 6))
+
+    def test_with_statement(self):
+        with ClientEncryption(
+                KMS_PROVIDERS, 'admin.datakeys',
+                client_context.client, OPTS) as client_encryption:
+            pass
+        with self.assertRaisesRegex(
+                InvalidOperation, 'Cannot use closed ClientEncryption'):
+            client_encryption.create_data_key('local')
+
+
 # Spec tests
 
 AWS_CREDS = {