Move encrypted router to tests

Author: aclark4life

django_mongodb_backend/routers.py

@@ -16,29 +16,3 @@ def allow_migrate(self, db, app_label, model_name=None, **hints):
         except LookupError:
             return None
         return False if issubclass(model, EmbeddedModel) else None
-
-
-class EncryptionRouter:
-    """
-    Routes database operations for 'encrypted' models to the 'encryption' DB.
-    """
-
-    def db_for_read(self, model, **hints):
-        if getattr(model, "encrypted_fields_map", False):
-            return "encryption"
-        return None
-
-    def db_for_write(self, model, **hints):
-        if getattr(model, "encrypted_fields_map", False):
-            return "encryption"
-        return None
-
-    def allow_migrate(self, db, app_label, model_name=None, **hints):
-        """
-        Ensure that the 'encrypted' models only appear in the 'encryption' DB,
-        and not in the default DB.
-        """
-        model = hints.get("model")
-        if model and getattr(model, "encrypted_fields_map", False):
-            return db == "encryption"
-        return None

tests/encryption_/routers.py

@@ -0,0 +1,26 @@
+from django.db import connection
+
+
+class EncryptedRouter:
+    """
+    Routes database operations for encrypted models to the encrypted DB.
+    """
+
+    def db_for_read(self, model, **hints):
+        with connection.schema_editor() as editor:
+            if model and getattr(editor._get_encrypted_fields_map(model), False):
+                return "encrypted"
+        return None
+
+    def db_for_write(self, model, **hints):
+        with connection.schema_editor() as editor:
+            if model and getattr(editor._get_encrypted_fields_map(model), False):
+                return "encrypted"
+        return None
+
+    def allow_migrate(self, db, app_label, model_name=None, **hints):
+        model = hints.get("model")
+        with connection.schema_editor() as editor:
+            if model and getattr(editor._get_encrypted_fields_map(model), False):
+                return db == "encrypted"
+        return None

tests/encryption_/tests.py

@@ -1,9 +1,11 @@
 from django.db import connection
-from django.test import TestCase
+from django.test import TestCase, override_settings
 
 from .models import Person
+from .routers import EncryptedRouter
 
 
+@override_settings(DATABASE_ROUTERS=[EncryptedRouter])
 class EncryptedModelTests(TestCase):
     @classmethod
     def setUpTestData(cls):