cleanup

Author: mhindery

djangosaml2/backends.py

@@ -16,6 +16,7 @@
 import logging
 from typing import Any, Optional, Tuple
 
+from django.apps import apps
 from django.conf import settings
 from django.contrib import auth
 from django.contrib.auth.backends import ModelBackend
@@ -27,30 +28,16 @@
 logger = logging.getLogger('djangosaml2')
 
 
-def get_model(model_path: str):
-    from django.apps import apps
-    try:
-        return apps.get_model(model_path)
-    except LookupError:
-        raise ImproperlyConfigured(f"SAML_USER_MODEL refers to model '{model_path}' that has not been installed")
-    except ValueError:
-        raise ImproperlyConfigured(f"SAML_USER_MODEL is {model_path}, but must be of the form 'app_label.model_name'")
-
-
 def set_attribute(obj: Any, attr: str, new_value: Any) -> bool:
     """ Set an attribute of an object to a specific value, if it wasn't that already.
         Return True if the attribute was changed and False otherwise.
     """
-
     if not hasattr(obj, attr):
         setattr(obj, attr, new_value)
         return True
-    else:
-        old_value = getattr(obj, attr)
-        if new_value != old_value:
-            setattr(obj, attr, new_value)
-            return True
-
+    if new_value != getattr(obj, attr):
+        setattr(obj, attr, new_value)
+        return True
     return False
 
 
@@ -64,7 +51,13 @@ class Saml2Backend(ModelBackend):
     def _user_model(self):
         """ Returns the user model specified in the settings, or the default one from this Django installation """
         if hasattr(settings, 'SAML_USER_MODEL'):
-            return get_model(settings.SAML_USER_MODEL)
+            try:
+                return apps.get_model(settings.SAML_USER_MODEL)
+            except LookupError:
+                raise ImproperlyConfigured(f"Model '{settings.SAML_USER_MODEL}' could not be loaded")
+            except ValueError:
+                raise ImproperlyConfigured(f"Model was specified as '{settings.SAML_USER_MODEL}', but it must be of the form 'app_label.model_name'")
+
         return auth.get_user_model()
 
     @property

tests/testprofiles/tests.py

@@ -21,29 +21,12 @@
 from django.core.exceptions import ImproperlyConfigured
 from django.test import TestCase, override_settings
 
-from djangosaml2.backends import (Saml2Backend, get_model, set_attribute)
+from djangosaml2.backends import Saml2Backend, set_attribute
 
 from .models import TestUser
 
-User = get_user_model()  # = TestUser
-
 
 class BackendUtilMethodsTests(TestCase):
-    def test_get_model_ok(self):
-        user_model = get_model('testprofiles.TestUser')
-        self.assertEqual(user_model, TestUser)
-
-    def test_get_model_nonexisting(self):
-        nonexisting_model = 'testprofiles.NonExisting'
-
-        with self.assertRaisesMessage(ImproperlyConfigured, f"SAML_USER_MODEL refers to model '{nonexisting_model}' that has not been installed"):
-            get_model(nonexisting_model)
-
-    def test_get_model_invalid_specifier(self):
-        nonexisting_model = 'random_package.specifier.testprofiles.NonExisting'
-
-        with self.assertRaisesMessage(ImproperlyConfigured, "SAML_USER_MODEL is random_package.specifier.testprofiles.NonExisting, but must be of the form 'app_label.model_name'"):
-            get_model(nonexisting_model)
 
     def test_set_attribute(self):
         u = TestUser()
@@ -74,6 +57,19 @@ def setUp(self):
         self.backend = self.backend_cls()
         self.user = TestUser.objects.create(username='john')
 
+    def test_get_model_ok(self):
+        self.assertEqual(self.backend._user_model, TestUser)
+
+    def test_get_model_nonexisting(self):
+        with override_settings(SAML_USER_MODEL='testprofiles.NonExisting'):
+            with self.assertRaisesMessage(ImproperlyConfigured, "Model 'testprofiles.NonExisting' could not be loaded"):
+                self.assertEqual(self.backend._user_model, None)
+
+    def test_get_model_invalid_specifier(self):
+        with override_settings(SAML_USER_MODEL='random_package.specifier.testprofiles.NonExisting'):
+            with self.assertRaisesMessage(ImproperlyConfigured, "Model was specified as 'random_package.specifier.testprofiles.NonExisting', but it must be of the form 'app_label.model_name'"):
+                self.assertEqual(self.backend._user_model, None)
+
     def test_user_model_specified(self):
         with override_settings(AUTH_USER_MODEL='auth.User'):
             with override_settings(SAML_USER_MODEL='testprofiles.TestUser'):
@@ -209,10 +205,10 @@ def test_create_user_with_required_fields(self):
         self.assertIs(user.email_verified, True)
 
     def test_django_user_main_attribute(self):
-        old_username_field = User.USERNAME_FIELD
-        User.USERNAME_FIELD = 'slug'
+        old_username_field = get_user_model().USERNAME_FIELD
+        get_user_model().USERNAME_FIELD = 'slug'
         self.assertEqual(self.backend._user_lookup_attribute, 'slug')
-        User.USERNAME_FIELD = old_username_field
+        get_user_model().USERNAME_FIELD = old_username_field
 
         with override_settings(AUTH_USER_MODEL='auth.User'):
             self.assertEqual(