fix encrypted tests

Author: timgraham

.github/workflows/mongodb_settings.py

@@ -1,5 +1,7 @@
 import os
 
+from pymongo.encryption_options import AutoEncryptionOpts
+
 from django_mongodb_backend import parse_uri
 
 if mongodb_uri := os.getenv("MONGODB_URI"):
@@ -27,6 +29,35 @@
         },
     }
 
+DATABASES["encrypted"] = {
+    "ENGINE": "django_mongodb_backend",
+    "NAME": "djangotests-encrypted",
+    "OPTIONS": {
+        "auto_encryption_opts": AutoEncryptionOpts(
+            key_vault_namespace="my_encrypted_database.keyvault",
+            kms_providers={"local": {"key": os.urandom(96)}},
+        )
+    },
+    "KMS_PROVIDERS": {},
+    "KMS_CREDENTIALS": {},
+}
+
+
+class EncryptedRouter:
+    def allow_migrate(self, db, app_label, model_name=None, **hints):
+        # The encryption_ app's models are only created in the encrypted database.
+        if app_label == "encryption_":
+            return db == "encrypted"
+        # Don't create other app's models in the encrypted database.
+        if db == "encrypted":
+            return False
+        return None
+
+    def kms_provider(self, model, **hints):
+        return "local"
+
+
+DATABASE_ROUTERS = [EncryptedRouter()]
 DEFAULT_AUTO_FIELD = "django_mongodb_backend.fields.ObjectIdAutoField"
 PASSWORD_HASHERS = ("django.contrib.auth.hashers.MD5PasswordHasher",)
 SECRET_KEY = "django_tests_secret_key"

tests/encryption_/routers.py

@@ -10,12 +10,12 @@ class TestEncryptedRouter:
 
     def allow_migrate(self, db, app_label, model_name=None, model=None, **hints):
         if model:
-            return db == ("other" if has_encrypted_fields(model) else "default")
+            return db == ("encrypted" if has_encrypted_fields(model) else "default")
         return db == "default"
 
     def db_for_read(self, model, **hints):
         if has_encrypted_fields(model):
-            return "other"
+            return "encrypted"
         return "default"
 
     db_for_write = db_for_read

tests/encryption_/tests.py

@@ -29,19 +29,19 @@
                 "bsonType": "string",
                 "path": "cc_type",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b" \x901\x89\x1f\xafAX\x9b*\xb1\xc7\xc5\xfdl\xa4", 4),
+                # "keyId": Binary(b" \x901\x89\x1f\xafAX\x9b*\xb1\xc7\xc5\xfdl\xa4", 4),
             },
             {
                 "bsonType": "long",
                 "path": "cc_number",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b"\x97\xb4\x9d\xb8\xd5\xa6Ay\x85\xfe\x00\xc0\xd4{\xa2\xff", 4),
+                # "keyId": Binary(b"\x97\xb4\x9d\xb8\xd5\xa6Ay\x85\xfe\x00\xc0\xd4{\xa2\xff", 4),
             },
             {
                 "bsonType": "decimal",
                 "path": "account_balance",
                 "queries": {"queryType": "range"},
-                "keyId": Binary(b"\xcc\x01-s\xea\xd9B\x8d\x80\xd7\xf8!n\xc6\xf5U", 4),
+                # "keyId": Binary(b"\xcc\x01-s\xea\xd9B\x8d\x80\xd7\xf8!n\xc6\xf5U", 4),
             },
         ]
     },
@@ -51,31 +51,31 @@
                 "bsonType": "string",
                 "path": "ssn",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b"\x14F\x89\xde\x8d\x04K7\xa9\x9a\xaf_\xca\x8a\xfb&", 4),
+                # "keyId": Binary(b"\x14F\x89\xde\x8d\x04K7\xa9\x9a\xaf_\xca\x8a\xfb&", 4),
             },
             {
                 "bsonType": "date",
                 "path": "birth_date",
                 "queries": {"queryType": "range"},
-                "keyId": Binary(b"@\xdd\xb4\xd2%\xc2B\x94\xb5\x07\xbc(ER[s", 4),
+                # "keyId": Binary(b"@\xdd\xb4\xd2%\xc2B\x94\xb5\x07\xbc(ER[s", 4),
             },
             {
                 "bsonType": "binData",
                 "path": "profile_picture",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b"Q\xa2\xebc!\xecD,\x8b\xe4$\xb6ul9\x9a", 4),
+                # "keyId": Binary(b"Q\xa2\xebc!\xecD,\x8b\xe4$\xb6ul9\x9a", 4),
             },
             {
                 "bsonType": "int",
                 "path": "patient_age",
                 "queries": {"queryType": "range"},
-                "keyId": Binary(b"\ro\x80\x1e\x8e1K\xde\xbc_\xc3bi\x95\xa6j", 4),
+                # "keyId": Binary(b"\ro\x80\x1e\x8e1K\xde\xbc_\xc3bi\x95\xa6j", 4),
             },
             {
                 "bsonType": "double",
                 "path": "weight",
                 "queries": {"queryType": "range"},
-                "keyId": Binary(b"\x9b\xfd:n\xe1\xd0N\xdd\xb3\xe7e)\x06\xea\x8a\x1d", 4),
+                # "keyId": Binary(b"\x9b\xfd:n\xe1\xd0N\xdd\xb3\xe7e)\x06\xea\x8a\x1d", 4),
             },
         ]
     },
@@ -85,33 +85,50 @@
                 "bsonType": "int",
                 "path": "patient_id",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b"\x8ft\x16:\x8a\x91D\xc7\x8a\xdf\xe5O\n[\xfd\\", 4),
+                # "keyId": Binary(b"\x8ft\x16:\x8a\x91D\xc7\x8a\xdf\xe5O\n[\xfd\\", 4),
             },
             {
                 "bsonType": "string",
                 "path": "patient_name",
-                "keyId": Binary(b"<\x9b\xba\xeb:\xa4@m\x93\x0e\x0c\xcaN\x03\xfb\x05", 4),
+                # "keyId": Binary(b"<\x9b\xba\xeb:\xa4@m\x93\x0e\x0c\xcaN\x03\xfb\x05", 4),
             },
             {
                 "bsonType": "string",
                 "path": "patient_notes",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b"\x01\xe7\xd1isnB$\xa9(gwO\xca\x10\xbd", 4),
+                # "keyId": Binary(b"\x01\xe7\xd1isnB$\xa9(gwO\xca\x10\xbd", 4),
             },
             {
                 "bsonType": "date",
                 "path": "registration_date",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b"F\xfb\xae\x82\xd5\x9a@\xee\xbfJ\xaf#\x9c:-I", 4),
+                # "keyId": Binary(b"F\xfb\xae\x82\xd5\x9a@\xee\xbfJ\xaf#\x9c:-I", 4),
             },
             {
                 "bsonType": "bool",
                 "path": "is_active",
                 "queries": {"queryType": "equality"},
-                "keyId": Binary(b"\xb2\xb5\xc4K53A\xda\xb9V\xa6\xa9\x97\x94\xea;", 4),
+                # "keyId": Binary(b"\xb2\xb5\xc4K53A\xda\xb9V\xa6\xa9\x97\x94\xea;", 4),
             },
+            {"bsonType": "string", "path": "email", "queries": {"queryType": "equality"}},
         ]
     },
+    "encryption__patientportaluser": {
+        "fields": [
+            {"bsonType": "string", "path": "ip_address", "queries": {"queryType": "equality"}},
+            {"bsonType": "string", "path": "url", "queries": {"queryType": "equality"}},
+        ]
+    },
+    "encryption__encryptednumbers": {
+        "fields": [
+            {"bsonType": "int", "path": "pos_bigint", "queries": {"queryType": "equality"}},
+            {"bsonType": "int", "path": "pos_smallint", "queries": {"queryType": "equality"}},
+            {"bsonType": "int", "path": "smallint", "queries": {"queryType": "equality"}},
+        ]
+    },
+    "encryption__appointment": {
+        "fields": [{"bsonType": "date", "path": "time", "queries": {"queryType": "equality"}}]
+    },
 }
 
 
@@ -144,7 +161,7 @@ def reload_module(module):
 )
 @override_settings(DATABASE_ROUTERS=[TestEncryptedRouter()])
 class EncryptedFieldTests(TransactionTestCase):
-    databases = {"default", "other"}
+    databases = {"default", "encrypted"}
     available_apps = ["django_mongodb_backend", "encryption_"]
 
     def setUp(self):
@@ -259,7 +276,7 @@ def test_get_encrypted_fields_map(self):
             },
         }
         self.maxDiff = None
-        with connections["other"].schema_editor() as editor:
+        with connections["encrypted"].schema_editor() as editor:
             db_table = self.patient._meta.db_table
             self.assertEqual(
                 editor._get_encrypted_fields_map(self.patient),
@@ -272,11 +289,18 @@ def test_show_schema_map(self):
         call_command(
             "showschemamap",
             "--database",
-            "other",
+            "encrypted",
             verbosity=0,
             stdout=out,
         )
-        self.assertEqual(json_util.dumps(EXPECTED_ENCRYPTED_FIELDS_MAP, indent=2), out.getvalue())
+        # Remove keyIds since they are different for each run.
+        output_json = json_util.loads(out.getvalue())
+        for table in output_json:
+            for field in output_json[table]["fields"]:
+                del field["keyId"]
+        # TODO: probably we don't need to test the entire mapping, otherwise it
+        # requires updates every time a new model or field is added!
+        self.assertEqual(EXPECTED_ENCRYPTED_FIELDS_MAP, output_json)
 
     def test_set_encrypted_fields_map_in_client(self):
         # TODO: Create new client with and without schema map provided then
@@ -315,16 +339,17 @@ def test_patientrecord(self):
         self.assertTrue(PatientRecord.objects.filter(weight__gte=175.0).exists())
 
         # Test encrypted patient record in unencrypted database.
-        conn_params = connections["other"].get_connection_params()
+        conn_params = connections["encrypted"].get_connection_params()
         if conn_params.pop("auto_encryption_opts", False):
             # Call MongoClient instead of get_new_connection because
             # get_new_connection will return the encrypted connection
             # from the connection pool.
-            connection = pymongo.MongoClient(**conn_params)
-            patientrecords = connection["test_other"].encryption__patientrecord.find()
-            ssn = patientrecords[0]["ssn"]
-            self.assertTrue(isinstance(ssn, Binary))
-            connection.close()
+            with pymongo.MongoClient(**conn_params) as new_connection:
+                patientrecords = new_connection[
+                    "test_django_encrypted"
+                ].encryption__patientrecord.find()
+                ssn = patientrecords[0]["ssn"]
+                self.assertTrue(isinstance(ssn, Binary))
 
     def test_patient(self):
         self.assertEqual(
@@ -343,9 +368,9 @@ def test_patient(self):
         )
 
         # Test decrypted patient record in encrypted database.
-        patients = connections["other"].database.encryption__patient.find()
+        patients = connections["encrypted"].database.encryption__patient.find()
         self.assertEqual(len(list(patients)), 1)
-        records = connections["other"].database.encryption__patientrecord.find()
+        records = connections["encrypted"].database.encryption__patientrecord.find()
         self.assertTrue("__safeContent__" in records[0])