Support custom user models in user conditions

This change adds support for custom user models to the user condition and its validator.

flags/conditions/conditions.py

@@ -2,6 +2,7 @@
 from distutils.util import strtobool
 
 import django
+from django.contrib.auth import get_user_model
 from django.utils import dateparse, timezone
 
 from flags.conditions.registry import register
@@ -33,7 +34,7 @@ def user_condition(username, request=None, **kwargs):
     if request is None:
         raise RequiredForCondition("request is required for condition 'user'")
 
-    return request.user.get_username() == username
+    return getattr(request.user, get_user_model().USERNAME_FIELD) == username
 
 
 @register("anonymous", validator=validate_boolean)

flags/conditions/validators.py

@@ -1,6 +1,7 @@
 import re
 from distutils.util import strtobool
 
+from django.contrib.auth import get_user_model
 from django.core.exceptions import ValidationError
 from django.core.validators import RegexValidator
 from django.utils import dateparse
@@ -36,11 +37,11 @@ def validate_boolean(value):
 
 
 def validate_user(value):
-    from django.contrib.auth.models import User
+    UserModel = get_user_model()
 
     try:
-        User.objects.get(username=value)
-    except User.DoesNotExist:
+        UserModel.objects.get(**{UserModel.USERNAME_FIELD: value})
+    except UserModel.DoesNotExist:
         raise ValidationError("Enter the username of a valid user.")
 
 

flags/tests/settings.py

@@ -32,7 +32,7 @@
     "debug_toolbar",
 )
 
-INSTALLED_APPS += ("flags",)
+INSTALLED_APPS += ("flags", "flags.tests.testapp")
 
 if django.VERSION >= (1, 10):  # pragma: no cover
     MIDDLEWARE = (

flags/tests/test_conditions_conditions.py

@@ -1,8 +1,9 @@
 from datetime import timedelta
 
+from django.contrib.auth import get_user_model
 from django.contrib.auth.models import AnonymousUser, User
 from django.http import HttpRequest, QueryDict
-from django.test import TestCase
+from django.test import TestCase, override_settings
 from django.utils import timezone
 
 from flags.conditions.conditions import (
@@ -15,7 +16,6 @@
     path_condition,
     user_condition,
 )
-from mock import MagicMock
 
 
 class BooleanConditionTestCase(TestCase):
@@ -62,12 +62,12 @@ def test_request_required(self):
         with self.assertRaises(RequiredForCondition):
             user_condition("testuser")
 
-    def test_with_custom_user(self):
-        mock_user = MagicMock()
-        mock_user.get_username.return_value = "test@test.com"
-        self.request.user = mock_user
-
-        self.assertTrue(user_condition("test@test.com", request=self.request))
+    @override_settings(AUTH_USER_MODEL="testapp.MyUserModel")
+    def test_custom_user_model_valid(self):
+        user = get_user_model()(identifier="customuser")
+        user.save()
+        self.request.user = user
+        self.assertTrue(user_condition("customuser", request=self.request))
 
 
 class AnonymousConditionTestCase(TestCase):

flags/tests/test_conditions_validators.py

@@ -1,6 +1,6 @@
-from django.contrib.auth.models import User
+from django.contrib.auth import get_user_model
 from django.core.exceptions import ValidationError
-from django.test import TestCase
+from django.test import TestCase, override_settings
 
 from flags.conditions.validators import (
     validate_boolean,
@@ -64,16 +64,27 @@ def test_valid_boolean_strings(self):
 
 
 class ValidateUserTestCase(TestCase):
-    def setUp(self):
-        User.objects.create_user(username="testuser", email="test@user")
-
     def test_invalid_user(self):
         with self.assertRaises(ValidationError):
             validate_user("nottestuser")
 
     def test_valid_user(self):
+        User = get_user_model()
+        User.objects.create_user(username="testuser", email="test@user")
         validate_user("testuser")
 
+    @override_settings(AUTH_USER_MODEL="testapp.MyUserModel")
+    def test_custom_user_valid(self):
+        User = get_user_model()
+        u = User(identifier="customuser")
+        u.save()
+        validate_user("customuser")
+
+    @override_settings(AUTH_USER_MODEL="testapp.MyUserModel")
+    def test_custom_user_invalid(self):
+        with self.assertRaises(ValidationError):
+            validate_user("nottestuser")
+
 
 class ValidateDateTestCase(TestCase):
     def test_invalid_date_strings(self):

flags/tests/testapp/migrations/0001_initial.py

@@ -0,0 +1,26 @@
+# Generated by Django 2.2.12 on 2020-05-06 08:26
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='MyUserModel',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('password', models.CharField(max_length=128, verbose_name='password')),
+                ('last_login', models.DateTimeField(blank=True, null=True, verbose_name='last login')),
+                ('identifier', models.CharField(max_length=40, unique=True)),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
+    ]

flags/tests/testapp/models.py

@@ -0,0 +1,8 @@
+from django.contrib.auth.models import AbstractBaseUser
+from django.db import models
+
+
+# Custom User model to test user conditions
+class MyUserModel(AbstractBaseUser):
+    identifier = models.CharField(max_length=40, unique=True)
+    USERNAME_FIELD = "identifier"

tox.ini

@@ -45,6 +45,7 @@ exclude=
     .tox,
     __pycache__,
     flags/migrations/*
+    flags/tests/testapp/migrations/*
 
 
 [isort]