Move encrypted_fields_map to schema (2/x)

- Refactor tests

Author: aclark4life

django_mongodb_backend/schema.py

@@ -421,20 +421,26 @@ def _field_should_have_unique(self, field):
 
     def _create_collection(self, model):
         """
-        Create a collection or encrypted collection for the model.
+        If the model is not encrypted, create a normal collection otherwise
+        create an encrypted collection with the encrypted fields map.
         """
 
-        if hasattr(model, "encrypted"):
+        if not hasattr(model, "encrypted"):
+            self.get_database().create_collection(model._meta.db_table)
+        else:
+            # TODO: Route to the encrypted database connection.
             auto_encryption_opts = self.connection.settings_dict.get("OPTIONS", {}).get(
                 "auto_encryption_opts"
             )
             client = self.connection.connection
+
             client_encryption = get_client_encryption(auto_encryption_opts, client)
             client_encryption.create_encrypted_collection(
                 client.database,
                 model._meta.db_table,
-                {"fields": []},
+                self._get_encrypted_fields_map(model),
                 "local",  # TODO: KMS provider should be configurable
             )
-        else:
-            self.get_database().create_collection(model._meta.db_table)
+
+    def _get_encrypted_fields_map(self, model):
+        return {"fields": []}

tests/encryption_/tests.py

@@ -1,32 +1,24 @@
+from django.db import connection
 from django.test import TestCase
 
 from .models import Person
 
 
 class EncryptedModelTests(TestCase):
-    databases = ["encryption"]
-
     @classmethod
     def setUpTestData(cls):
-        cls.objs = [Person.objects.create()]
-
-    def test_encrypted_fields_map_on_class(self):
-        expected = {
-            "fields": {
-                "ssn": "EncryptedCharField",
-            }
-        }
-        self.assertEqual(Person.encrypted_fields_map, expected)
+        cls.person = Person(ssn="123-45-6789")
 
     def test_encrypted_fields_map_on_instance(self):
-        instance = Person(ssn="123-45-6789")
         expected = {
             "fields": {
                 "ssn": "EncryptedCharField",
             }
         }
-        self.assertEqual(instance.encrypted_fields_map, expected)
+        with connection.schema_editor() as editor:
+            self.assertEqual(editor._get_encrypted_fields_map(self.person), expected)
 
     def test_non_encrypted_fields_not_included(self):
-        encrypted_field_names = Person.encrypted_fields_map.get("fields").keys()
-        self.assertNotIn("name", encrypted_field_names)
+        with connection.schema_editor() as editor:
+            encrypted_field_names = editor._get_encrypted_fields_map(self.person).get("fields")
+            self.assertNotIn("name", encrypted_field_names)