diff options
Diffstat (limited to '')
-rw-r--r-- | ctrack/users/models.py | 10 | ||||
-rw-r--r-- | ctrack/users/tests/test_models.py | 3 |
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 |