further asserts added to the test

2 files changed, 14 insertions, 1 deletions

diff --git a/ctrack/organisations/models.py b/ctrack/organisations/models.py
index 9eb0a23..78e4801 100644
--- a/ctrack/organisations/models.py
+++ b/ctrack/organisations/models.py
@@ -155,6 +155,13 @@ class Organisation(models.Model):
     def applicable_systems(self):
         return self.applicablesystem_set.all()
 
+    def systems(self):
+        ess = self.essentialservice_set.all()
+        out = []
+        for es in ess:
+            out.extend(list(es.systems.all()))
+        return out
+
 
 class Address(models.Model):
     organisation = models.ForeignKey(
diff --git a/ctrack/organisations/tests/test_models.py b/ctrack/organisations/tests/test_models.py
index 09d4f12..103ead4 100644
--- a/ctrack/organisations/tests/test_models.py
+++ b/ctrack/organisations/tests/test_models.py
@@ -44,8 +44,14 @@ def test_essential_service(org):
     ass = ApplicableSystemFactory.create(
         name=random.choice(fnames), organisation=org, caf=caf,
     )
+    ass2 = ApplicableSystemFactory.create(
+        name=random.choice(fnames), organisation=org, caf=caf,
+    )
     es = EssentialService.objects.create(
         name="Test ES", description="Test ES Description", organisation=org
     )
-    es.systems.add(ass)
+    es.systems.add(ass, ass2)
     assert es.systems.first().organisation.name == org.name
+    assert es.name == "Test ES"
+    assert es.systems.count() == 2
+    assert ass.name in [s.name for s in org.systems()]