many improvements

test/rota_constraints/test_avoid_grades_by_week.py

@@ -0,0 +1,195 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days
+
+from rota.workers import Worker
+
+import itertools
+
+
+def generate_basic_rota(weeks_to_rota=9):
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    # Add a few workers
+    Rota.add_workers(
+        [
+            Worker(name="worker1", site="group1", grade=1),
+            Worker(name="worker2", site="group1", grade=2),
+            Worker(name="worker3", site="group1", grade=3),
+        ]
+    )
+
+    return Rota
+
+
+def date_generator(from_date, days):
+    n = 0
+    while True:
+        yield from_date
+
+        n = n + 1
+        if n >= days:
+            break
+        from_date = from_date + datetime.timedelta(days=1)
+
+
+class TestWorkerRequests:
+    def test_avoid_grade_constraint(self):
+        Rota = generate_basic_rota()
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1], [1,2,3], ["a"])
+        Rota.add_grade_constraint_by_week([2], [4,5,6], ["a"])
+        Rota.add_grade_constraint_by_week([3], [7,8,9], ["a"])
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("avoid_grade_constraint")
+
+        shift_summary = Rota.get_shift_summary_dict()
+        for worker_name in shift_summary:
+            worker_shifts = shift_summary[worker_name]
+
+            worker = Rota.get_worker_by_name(worker_name)
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            match worker_name:
+                case "worker1":
+                    assert shift_string.startswith("-"*21)
+                case "worker2":
+                    assert shift_string.startswith(21*"a"+"-"*21)
+                case "worker3":
+                    assert shift_string.endswith("-"*21)
+
+
+    def test_avoid_grade_constraint_invalid_shift(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+        with pytest.raises(ValueError):
+
+            Rota.add_grade_constraint_by_week([1], [1,2,3], ["b"])
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="b", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1], [1,2,3], ["b"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+
+    def test_avoid_grade_constraint_invalid_grade(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+        with pytest.raises(ValueError):
+
+            Rota.add_grade_constraint_by_week([4], [1,2,3], ["a"])
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+        Rota.add_worker(
+                Worker(name="worker4", site="group1", grade=4),
+        )
+
+
+        Rota.add_grade_constraint_by_week([4], [1,2,3], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+    def test_avoid_grade_constraint_multiple_grades(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1,2], [1,2,3], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+        Rota.add_grade_constraint_by_week([3], [1,2,3], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "warning"
+
+    def test_avoid_grade_constraint_multiple_shifts(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+            SingleShift(
+                sites=("group1",), name="b", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1], [1,2,3], ["a", "b"])
+        Rota.add_grade_constraint_by_week([2], [4,5,6], ["a", "b"])
+        Rota.add_grade_constraint_by_week([3], [7,8,9], ["a", "b"])
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("avoid_grade_constraint")
+
+        shift_summary = Rota.get_shift_summary_dict()
+        for worker_name in shift_summary:
+            worker_shifts = shift_summary[worker_name]
+
+            worker = Rota.get_worker_by_name(worker_name)
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            match worker_name:
+                case "worker1":
+                    assert shift_string.startswith("-"*21)
+                case "worker2":
+                    assert "-"*21 in shift_string # hacky
+                case "worker3":
+                    assert shift_string.endswith("-"*21)
+
+        Rota.add_grade_constraint_by_week([2], [7,8,9], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+        Rota.add_grade_constraint_by_week([1], [7,8,9], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "warning"