Merge pull request #106 from defencedigital/postgres-migration

Adds new Engagement Strategy stuff

1 files changed, 19 insertions, 18 deletions

diff --git a/instruments/tests/test_models.py b/instruments/tests/test_models.py
index 0d7eafa..7762dde 100644
--- a/instruments/tests/test_models.py
+++ b/instruments/tests/test_models.py
@@ -1,29 +1,30 @@
 import datetime
 
 import pytest
-from django.test import TestCase
 
 from engagements import models
 from engagements.utils import populate_database
 
+pytestmark = pytest.mark.django_db
 
-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.fixture
+def test_data():
+    return populate_database()
 
-    @pytest.mark.django_db
-    def test_get_hours_of_effort_for_dsc_for_org(self):
-        org = self.data["orgs"][0]
+
+class TestModels:
+    def test_check_all_dcs(self, test_data):
+        dscs = test_data.get("sub_instruments")
+        assert dscs[0].title == "DSC 1 - Title 1"
+
+    def test_get_hours_of_effort_for_dsc_for_org(self, test_data):
+        org = test_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]
+        si = test_data["sub_instruments"][0]
+        si2 = test_data["sub_instruments"][2]
+        si3 = test_data["sub_instruments"][3]
+        # si_not = test_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),
@@ -41,6 +42,6 @@ class TestModels(TestCase):
         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)
+        assert si.effort_for_org(org) == 2
+        assert si2.effort_for_org(org) == 2
+        assert si3.effort_for_org(org) == 2