DEV Community

Gahyun Son
Gahyun Son

Posted on

unittest with get_user_model()

1. Custom User model

In Django, you can override the user model provided by Django.

  1. AbstractUser. 2. AbstractBaseUser.

I chose AbstractBaseUser. Below code is a my custom user.

I defined in /djangoproject/app/core/models.py
(I'm going to write all models on this file.)

from django.db import models

from django.contrib.auth.models import (
    AbstractBaseUser,
    BaseUserManager,
    PermissionsMixin,
)


class User(AbstractBaseUser, PermissionsMixin):
    name = models.CharField(max_length=255)
    email = models.EmailField(max_length=255, unique=True)
    company = models.CharField(max_length=100, blank=True, null=True)
    address = models.TextField(max_length=200, blank=True, null=True)
    address2 = models.TextField(max_length=200, blank=True, null=True)
    city = models.CharField(max_length=50, blank=True, null=True)
    nation = models.CharField(max_length=100, blank=True, null=True)
    postal = models.CharField(max_length=30, blank=True, null=True)
    phone = models.CharField(max_length=16, blank=True, null=True)

    is_active = models.BooleanField(default=True)
    is_staff = models.BooleanField(default=False)

    objects = UserManager()

    USERNAME_FIELD = 'email'
Enter fullscreen mode Exit fullscreen mode

I made UserManager too, but I haven't included it here for brevity.

  1. Setting AUTH_USER_MODEL in settings.py
AUTH_USER_MODEL = 'core.user'
Enter fullscreen mode Exit fullscreen mode

Add it to settings.py. The value is a my custom user model.

  1. Creating unittest a file. /djangoproject/app/core/tests/test_model.py
class ModelTests(TestCase):
    def test_create_user_with_email_successful(self):
        email = 'test@example.com'
        password = 'test123test'
        user = get_user_model().objects.create_user(
            email=email,
            password=password
        )

        self.assertEqual(user.email, email)
        self.assertTrue(user.check_password(password))
Enter fullscreen mode Exit fullscreen mode
  • I checked whether the email is correctly saved in the database and whether the password is properly validated.

  • There are two ways to reference the custom user model.

  • Direct import user model.

from core import models

User = models.User
user = User.objects.create_user(...)
Enter fullscreen mode Exit fullscreen mode

It works, but if the AUTH_USER_MODEL setting is later changed to a different user model, you'll need to update all instances to the changed model.

  1. Import get_user_model().
from django.contrib.auth import get_user_model

user = get_user_model().objects.create_user(...)
Enter fullscreen mode Exit fullscreen mode

get_user_model returns the currently active user model. It will work without requiring code modifications.

Top comments (0)