more testing for stakeholder

2 files changed, 13 insertions, 0 deletions

diff --git a/ctrack/users/models.py b/ctrack/users/models.py
index b53c217..688254f 100644
--- a/ctrack/users/models.py
+++ b/ctrack/users/models.py
@@ -15,3 +15,13 @@ class User(AbstractUser):
 
     def get_absolute_url(self):
         return reverse("users:detail", kwargs={"username": self.username})
+
+    def is_stakeholder(self):
+        if self.stakeholder is not None:
+            return True
+        else:
+            return False
+
+    def get_organisation_name(self):
+        if self.is_stakeholder():
+            return self.stakeholder.person.organisation.name
diff --git a/ctrack/users/tests/test_models.py b/ctrack/users/tests/test_models.py
index 504195f..402c41b 100644
--- a/ctrack/users/tests/test_models.py
+++ b/ctrack/users/tests/test_models.py
@@ -22,5 +22,8 @@ def test_stakeholder_model(person, user):
     the system.
     """
     stakeholder = Stakeholder(person=person)
+    org = person.organisation.name
     user.stakeholder = stakeholder
     assert user.stakeholder.person.first_name == "Chinaplate"
+    assert user.is_stakeholder() is True
+    assert user.get_organisation_name() == org