fix tests

Author: JBKahn

django_sharding_library/router.py

@@ -74,22 +74,23 @@ def allow_migrate(self, db, app_label, model_name=None, **hints):
             return False
         model_name = model_name or hints.get('model_name')
         model = hints.get('model')
-        if not model_name and not model:
+        if model:
+            model_name = model.__name__
+        if not model_name:
             raise InvalidMigrationException(
                 'Model name not provided in migration, please pass a `model_name` or `model` with the hints passed into the migration.'
             )
 
         # Sometimes, when extending models from another app i.e. the User Model, the app label
         # is the app label of the app where the change is defined but to app with the model is
         # passed in with the model name.
-        if not model:
-            try:
-                app = apps.get_app_config(app_label)
-                model = app.get_model(model_name)
-            except LookupError:
-                app_label = model_name.split('.')[0]
-                app = apps.get_app_config(app_label)
-                model = app.get_model(model_name[len(app_label) + 1:])
+        try:
+	    app = apps.get_app_config(app_label)
+            model = app.get_model(model_name)
+        except LookupError:
+	    app_label = model_name.split('.')[0]
+	    app = apps.get_app_config(app_label)
+            model = app.get_model(model_name[len(app_label) + 1:])
 
         single_database = self.get_specific_database_or_none(model)
         shard_group = self.get_shard_group_if_sharded_or_none(model)

tests/test_router.py

@@ -169,6 +169,7 @@ def test_requires_model_name_to_be_passed_in(self):
         self.sut.allow_migrate(db='default', app_label='tests', **hints)
 
     def test_requires_model_to_be_passed_in(self):
+        from django.contrib.auth import get_user_model
         with self.assertRaises(InvalidMigrationException):
             self.sut.allow_migrate(db='default', app_label='tests')