Check for encrypted fields from unencrypted conn

Re-adding test removed in 46ca9dc as
assertEncrypted class method.

Author: aclark4life

tests/encryption_/test_fields.py

@@ -1,6 +1,12 @@
 import datetime
 from decimal import Decimal
 
+import pymongo
+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
 
 from .models import (
@@ -32,22 +38,69 @@
 from .test_base import EncryptionTestCase
 
 
-class EncryptedEmbeddedModelTests(EncryptionTestCase):
+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()
+        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.")
+
+
+class EncryptedEmbeddedModelTests(EncryptedFieldTests):
     def setUp(self):
         self.billing = Billing(cc_type="Visa", cc_number="4111111111111111")
         self.patient_record = PatientRecord(ssn="123-45-6789", billing=self.billing)
         self.patient = Patient.objects.create(
             patient_name="John Doe", patient_id=123456789, patient_record=self.patient_record
         )
 
-    def test_patient(self):
+    def test_object(self):
         patient = Patient.objects.get(id=self.patient.id)
         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")
 
 
-class EncryptedEmbeddedModelArrayTests(EncryptionTestCase):
+class EncryptedEmbeddedModelArrayTests(EncryptedFieldTests):
     def setUp(self):
         self.actor1 = Actor(name="Actor One")
         self.actor2 = Actor(name="Actor Two")
@@ -56,13 +109,14 @@ def setUp(self):
             cast=[self.actor1, self.actor2],
         )
 
-    def test_movie_actors(self):
+    def test_array(self):
         self.assertEqual(len(self.movie.cast), 2)
         self.assertEqual(self.movie.cast[0].name, "Actor One")
         self.assertEqual(self.movie.cast[1].name, "Actor Two")
+        self.assertEncrypted(self.movie, "cast")
 
 
-class EncryptedFieldTests(EncryptionTestCase):
+class EncryptedFieldTests(EncryptedFieldTests):
     def assertEquality(self, model_cls, val):
         model_cls.objects.create(value=val)
         fetched = model_cls.objects.get(value=val)
@@ -80,28 +134,36 @@ def assertRange(self, model_cls, *, low, high, threshold):
     # Equality-only fields
     def test_binary(self):
         self.assertEquality(BinaryModel, b"\x00\x01\x02")
+        self.assertEncrypted(BinaryModel, "value")
 
     def test_boolean(self):
         self.assertEquality(BooleanModel, True)
+        self.assertEncrypted(BooleanModel, "value")
 
     def test_char(self):
         self.assertEquality(CharModel, "hello")
+        self.assertEncrypted(CharModel, "value")
 
     def test_email(self):
         self.assertEquality(EmailModel, "[email protected]")
+        self.assertEncrypted(EmailModel, "value")
 
     def test_ip(self):
         self.assertEquality(GenericIPAddressModel, "192.168.0.1")
+        self.assertEncrypted(GenericIPAddressModel, "value")
 
     def test_text(self):
         self.assertEquality(TextModel, "some text")
+        self.assertEncrypted(TextModel, "value")
 
     def test_url(self):
         self.assertEquality(URLModel, "https://example.com")
+        self.assertEncrypted(URLModel, "value")
 
     # Range fields
     def test_big_integer(self):
         self.assertRange(BigIntegerModel, low=100, high=200, threshold=150)
+        self.assertEncrypted(BigIntegerModel, "value")
 
     def test_date(self):
         self.assertRange(
@@ -110,6 +172,7 @@ def test_date(self):
             high=datetime.date(2024, 6, 10),
             threshold=datetime.date(2024, 6, 5),
         )
+        self.assertEncrypted(DateModel, "value")
 
     def test_datetime(self):
         self.assertRange(
@@ -118,6 +181,7 @@ def test_datetime(self):
             high=datetime.datetime(2024, 6, 2, 12, 0),
             threshold=datetime.datetime(2024, 6, 2, 0, 0),
         )
+        self.assertEncrypted(DateTimeModel, "value")
 
     def test_decimal(self):
         self.assertRange(
@@ -126,6 +190,7 @@ def test_decimal(self):
             high=Decimal("200.50"),
             threshold=Decimal("150"),
         )
+        self.assertEncrypted(DecimalModel, "value")
 
     def test_duration(self):
         self.assertRange(
@@ -134,24 +199,31 @@ def test_duration(self):
             high=datetime.timedelta(days=10),
             threshold=datetime.timedelta(days=5),
         )
+        self.assertEncrypted(DurationModel, "value")
 
     def test_float(self):
         self.assertRange(FloatModel, low=1.23, high=4.56, threshold=3.0)
+        self.assertEncrypted(FloatModel, "value")
 
     def test_integer(self):
         self.assertRange(IntegerModel, low=5, high=10, threshold=7)
+        self.assertEncrypted(IntegerModel, "value")
 
     def test_positive_big_integer(self):
         self.assertRange(PositiveBigIntegerModel, low=100, high=500, threshold=200)
+        self.assertEncrypted(PositiveBigIntegerModel, "value")
 
     def test_positive_integer(self):
         self.assertRange(PositiveIntegerModel, low=10, high=20, threshold=15)
+        self.assertEncrypted(PositiveIntegerModel, "value")
 
     def test_positive_small_integer(self):
         self.assertRange(PositiveSmallIntegerModel, low=5, high=8, threshold=6)
+        self.assertEncrypted(PositiveSmallIntegerModel, "value")
 
     def test_small_integer(self):
         self.assertRange(SmallIntegerModel, low=-5, high=2, threshold=0)
+        self.assertEncrypted(SmallIntegerModel, "value")
 
     def test_time(self):
         self.assertRange(
@@ -160,9 +232,10 @@ def test_time(self):
             high=datetime.time(15, 0),
             threshold=datetime.time(12, 0),
         )
+        self.assertEncrypted(TimeModel, "value")
 
 
-class EncryptedFieldMixinTests(EncryptionTestCase):
+class EncryptedFieldMixinTests(EncryptedFieldTests):
     def test_null_true_raises_error(self):
         with self.assertRaisesMessage(
             ValueError, "'null=True' is not supported for encrypted fields."