Fix assertEncrypted for all but nested fields

Author: aclark4life

tests/encryption_/test_fields.py

@@ -5,7 +5,6 @@
 from bson.binary import Binary
 from django.conf import settings
 from django.db import connections
-from django.db.models import Model
 
 from django_mongodb_backend.fields import EncryptedCharField
 
@@ -39,50 +38,15 @@
 
 
 class EncryptedFieldTests(EncryptionTestCase):
-    def assertEncrypted(self, model_or_instance, field_name):
-        """
-        Check if the field value in the database is stored as Binary.
-        Works with either a Django model instance or a model class.
-        """
-
-        conn_params = connections["encrypted"].get_connection_params()
+    def assertEncrypted(self, model, field):
+        # Unencrypted connection to encrypted database
+        conn_params = connections["default"].get_connection_params()
         db_name = settings.DATABASES["encrypted"]["NAME"]
-
-        if conn_params.pop("auto_encryption_opts", False):
-            with pymongo.MongoClient(**conn_params) as new_connection:
-                if hasattr(model_or_instance, "_meta"):
-                    collection_name = model_or_instance._meta.db_table
-                else:
-                    self.fail(f"Object {model_or_instance!r} is not a Django model or instance")
-
-                collection = new_connection[db_name][collection_name]
-
-                # If it's an instance of a Django model, narrow to that _id
-                if isinstance(model_or_instance, Model):
-                    docs = collection.find(
-                        {"_id": model_or_instance.pk, field_name: {"$exists": True}}
-                    )
-                else:
-                    # Otherwise it's a model class
-                    docs = collection.find({field_name: {"$exists": True}})
-
-                found = False
-                for doc in docs:
-                    found = True
-                    value = doc.get(field_name)
-                    self.assertTrue(
-                        isinstance(value, Binary),
-                        msg=f"Field '{field_name}' in document {doc['_id']} is "
-                        "not encrypted (type={type(value)})",
-                    )
-
-                self.assertTrue(
-                    found,
-                    msg=f"No documents with field '{field_name}' found in '{{collection_name}}'",
-                )
-
-        else:
-            self.fail("auto_encryption_opts is not configured; encryption not enabled.")
+        with pymongo.MongoClient(**conn_params) as new_connection:
+            db = new_connection[db_name]
+            collection_name = model._meta.db_table
+            data = db[collection_name].find()[0]
+            self.assertIsInstance(data[field], Binary)
 
 
 class EncryptedEmbeddedModelTests(EncryptedFieldTests):
@@ -98,6 +62,7 @@ def test_object(self):
         self.assertEqual(patient.patient_record.ssn, "123-45-6789")
         self.assertEqual(patient.patient_record.billing.cc_type, "Visa")
         self.assertEqual(patient.patient_record.billing.cc_number, "4111111111111111")
+        # self.assertEncrypted(patient, ["patient_record", "ssn"])
 
 
 class EncryptedEmbeddedModelArrayTests(EncryptedFieldTests):