From c6b9f6e8b657096a63d65fcd54820b3007269bbe Mon Sep 17 00:00:00 2001
From: Cody Goodman <cody@rewardstyle.com>
Date: Tue, 22 Apr 2014 11:56:09 -0500
Subject: [PATCH] Add support for custom user models

---
 tos/models.py            | 3 ++-
 tos/tests/test_models.py | 4 +++-
 tos/tests/test_views.py  | 4 +++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/tos/models.py b/tos/models.py
index d1ac4fb..0468008 100644
--- a/tos/models.py
+++ b/tos/models.py
@@ -1,8 +1,9 @@
 from django.core.exceptions import ValidationError
 from django.db import models
-from django.contrib.auth.models import User
+from django.contrib.auth import get_user_model
 from django.utils.translation import ugettext_lazy as _
 
+User = get_user_model()
 
 class NoActiveTermsOfService(ValidationError):
     pass
diff --git a/tos/tests/test_models.py b/tos/tests/test_models.py
index 5b99486..1e94948 100644
--- a/tos/tests/test_models.py
+++ b/tos/tests/test_models.py
@@ -1,8 +1,10 @@
-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.db.models.signals import post_syncdb 
 
+User = get_user_model()
+
 from tos.models import TermsOfService, UserAgreement, has_user_agreed_latest_tos
 
 class TestModels(TestCase):
diff --git a/tos/tests/test_views.py b/tos/tests/test_views.py
index 1cf10c5..a8411e1 100644
--- a/tos/tests/test_views.py
+++ b/tos/tests/test_views.py
@@ -1,11 +1,13 @@
 from django.conf import settings
-from django.contrib.auth.models import User 
+from django.contrib.auth import get_user_model
 from django.core.exceptions import ValidationError 
 from django.core.urlresolvers import reverse 
 from django.test import TestCase 
 
 from tos.models import TermsOfService, UserAgreement, has_user_agreed_latest_tos
 
+User = get_user_model()
+
 class TestViews(TestCase):
 
     def setUp(self):