major test refactoring

test/test_shift_constraints.py

@@ -0,0 +1,71 @@
+import pytest
+import datetime
+from copy import deepcopy
+from rota.shifts import RotaBuilder, SingleShift, days
+from rota.workers import Worker
+
+@pytest.fixture
+def rota_2workers():
+    weeks_to_rota = 10
+    start_date = datetime.date(2022, 3, 7)
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+    )
+    worker1 = Worker(name="worker1", site="group1", grade=1)
+    worker2 = Worker(name="worker2", site="group1", grade=1)
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=1,
+        ),
+    )
+    return Rota
+
+@pytest.fixture
+def rota_3workers(rota_2workers):
+    Rota3 = deepcopy(rota_2workers)
+    worker3 = Worker(name="worker3", site="group1", grade=1)
+    Rota3.add_worker(worker3)
+    return Rota3
+
+def test_max_shifts(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_month"] = 14
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status == "ok"
+    assert rota_2workers.results.solver.termination_condition == "optimal"
+
+def test_max_shifts_fail(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_month"] = 13
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status in ("warning", "error")
+    assert rota_2workers.results.solver.termination_condition == "infeasible"
+
+def test_max_shifts_extra_worker(rota_3workers):
+    rota_3workers.constraint_options["max_shifts_per_month"] = 12
+    rota_3workers.build_and_solve()
+    assert rota_3workers.results.solver.status == "ok"
+    assert rota_3workers.results.solver.termination_condition == "optimal"
+
+def test_max_shifts_per_week_fail(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_week"] = 3
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status in ("warning", "error")
+    assert rota_2workers.results.solver.termination_condition == "infeasible"
+
+def test_max_shifts_per_week_pass(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_month"] = 14
+    rota_2workers.constraint_options["max_shifts_per_week"] = 4
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status == "ok"
+
+def test_max_shifts_per_week_extra_worker_pass(rota_3workers):
+    rota_3workers.constraint_options["max_shifts_per_week"] = 3
+    rota_3workers.build_and_solve()
+    assert rota_3workers.results.solver.status == "ok"