1 files changed, 46 insertions, 0 deletions
diff --git a/instruments/tests/test_models.py b/instruments/tests/test_models.py
new file mode 100644
index 0000000..0d7eafa
--- /dev/null
+++ b/instruments/tests/test_models.py
@@ -0,0 +1,46 @@
+import datetime
+
+import pytest
+from django.test import TestCase
+
+from engagements import models
+from engagements.utils import populate_database
+
+
+class TestModels(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.data = populate_database()
+
+    def test_check_all_dcs(self):
+        dscs = self.data.get("sub_instruments")
+        self.assertEqual(dscs[0].title, "DSC 1 - Title 1")
+
+    @pytest.mark.django_db
+    def test_get_hours_of_effort_for_dsc_for_org(self):
+        org = self.data["orgs"][0]
+        et = models.EngagementType.objects.get(name="INSPECTION")
+        si = self.data["sub_instruments"][0]
+        si2 = self.data["sub_instruments"][2]
+        si3 = self.data["sub_instruments"][3]
+        # si_not = self.data["sub_instruments"][1]
+        engagement = models.Engagement.objects.create(
+            proposed_start_date=datetime.date(2022, 10, 10),
+            proposed_end_date=datetime.date(2022, 10, 10),
+            engagement_type=et,
+            external_party=org,
+        )
+        ef1 = models.EngagementEffort.objects.create(
+            is_planned=True,
+            effort_type="REGULATION",
+            proposed_start_date=datetime.datetime(2022, 10, 10, 10, 0),
+            proposed_end_date=datetime.datetime(2022, 10, 10, 12, 0),
+            engagement=engagement,
+        )
+        ef1.sub_instruments.add(si)  # DSC 1
+        ef1.sub_instruments.add(si2)  # DSC 3
+        ef1.sub_instruments.add(si3)  # DSC 4
+        ef1.save()
+        self.assertEqual(si.effort_for_org(org), 2)
+        self.assertEqual(si2.effort_for_org(org), 2)
+        self.assertEqual(si3.effort_for_org(org), 2)