2 files changed, 42 insertions, 6 deletions
diff --git a/ctrack/users/tests/test_functional.py b/ctrack/users/tests/test_functional.py
index c68382c..d5390f6 100644
--- a/ctrack/users/tests/test_functional.py
+++ b/ctrack/users/tests/test_functional.py
@@ -16,9 +16,31 @@ from ctrack.users.models import User
 pytestmark = pytest.mark.django_db
 
 
-def test_user_can_log_in(browser, person, live_server):
+def test_regular_user_can_log_in(browser, person, live_server):
 
     # Toss McBride is an OES user. He logs into the system...
+    User.objects.create_user(username="toss", password="knob")
+    browser.get(live_server + "/accounts/login")
+    browser.find_element_by_id("id_login").send_keys("toss")
+    browser.find_element_by_id("id_password").send_keys("knob")
+    browser.find_element_by_id("sign_in_button").submit()
+    time.sleep(1)
+    current_url = browser.current_url
+    assert current_url == live_server + "/"
+
+    # On the other side, he sees some basic details about himself.
+    assert "ctrack - Department for Transport" in browser.title
+
+    h1 = browser.find_element_by_tag_name("h1")
+    assert h1.text == "Welcome to ctrack - Department for Transport"
+    type_user_message = browser.find_elements_by_tag_name("p")
+    assert "THIS IS A TEMPLATE FOR A REGULAR USER" in [
+        m.text for m in type_user_message
+    ]
+
+
+def test_stakeholder_can_log_in_and_see_their_home(browser, person, live_server):
+    # Toss McBride is an OES user. He logs into the system...
     stakeholder = Stakeholder.objects.create(person=person)
 
     user = User.objects.create_user(username="toss", password="knob")
@@ -39,6 +61,6 @@ def test_user_can_log_in(browser, person, live_server):
     h1 = browser.find_element_by_tag_name("h1")
     assert h1.text == "Welcome to ctrack - Department for Transport"
     type_user_message = browser.find_elements_by_tag_name("p")
-    assert "THIS IS A TEMPLATE FOR A REGULAR USER" in [
+    assert "THIS IS A TEMPLATE FOR A STAKEHOLDER USER" in [
         m.text for m in type_user_message
     ]
diff --git a/ctrack/users/tests/test_views.py b/ctrack/users/tests/test_views.py
index 6859215..0ba8b8e 100644
--- a/ctrack/users/tests/test_views.py
+++ b/ctrack/users/tests/test_views.py
@@ -1,8 +1,5 @@
 import pytest
-from django.contrib.auth import get_user_model
-from django.http import HttpRequest
 from django.test import RequestFactory
-from django.urls import resolve
 
 from ctrack.core.views import home_page
 from ctrack.organisations.models import Stakeholder
@@ -98,7 +95,7 @@ def test_home_page_h1_tag_with_client(client, django_user_model):
     assert b"</html>" in response.content
 
 
-def test_regular_user_gets_regular_user_template(django_user_model):
+def test_regular_user_redirected_to_their_template_on_login(django_user_model):
     """
     When a user logs in without a stakeholder mapping, they get sent to the regular user
     template.
@@ -110,3 +107,20 @@ def test_regular_user_gets_regular_user_template(django_user_model):
     response = home_page(request)
     assert response.status_code == 200
     assert b"<p>THIS IS A TEMPLATE FOR A REGULAR USER</p>" in response.content
+
+
+def test_stakeholder_redirected_to_their_template_on_login(django_user_model, person):
+    """
+    When a user logs in WITH a stakeholder mapping, they get sent to the stakehoder user
+    template.
+    """
+    user = django_user_model.objects.create_user(username="toss", password="knob")
+    stakeholder = Stakeholder.objects.create(person=person)
+    user.stakeholder = stakeholder
+    user.save()
+    factory = RequestFactory()
+    request = factory.get("/")
+    request.user = user
+    response = home_page(request)
+    assert response.status_code == 200
+    assert b"THIS IS A TEMPLATE FOR A STAKEHOLDER USER" in response.content