fix a few more things

test/rota_constraints/test_summed_shift_grade_constraints.py

@@ -0,0 +1,94 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days, MinSummedGradeByShiftsPerDayConstraint
+
+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),
+            Worker(name="worker4", site="group1", grade=1),
+            Worker(name="worker5", site="group1", grade=2),
+            Worker(name="worker6", site="group1", grade=3),
+        ]
+    )
+
+    return Rota
+
+
+def test_summed_shift_grade_constraints():
+    Rota = generate_basic_rota(5)
+
+    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,
+        ),
+    )
+
+    # Add summed shift grade constraints
+    Rota.constraint_options["min_summed_grade_by_shifts_per_day"].append(
+        MinSummedGradeByShiftsPerDayConstraint(
+            shifts=["a", "b"],
+            min_grade_sum=4,
+        ),
+    )
+
+    Rota.build_and_solve(options={"ratio": 0.000})
+
+    assert Rota.results.solver.status == "ok"
+
+    Rota.export_rota_to_html("summed_shift_grade_constraints", folder="tests")
+
+    a_grades = []
+    b_grades = []
+
+    for shift, shift_timetable in Rota.get_shift_timetable_by_week().items():
+        for week, week_timetable in shift_timetable.items():
+            for day, workers in week_timetable.items():
+                if shift == "a":
+                    for worker in workers:
+                        a_grades.append(worker.grade)
+                elif shift == "b":
+                    for worker in workers:
+                        b_grades.append(worker.grade)
+
+    # Check that the minimum summed grade is respected
+    summed_grades = [a + b for a, b in zip(a_grades, b_grades)]
+    for summed_grade in summed_grades:
+        assert summed_grade >= 4, f"Summed grade {summed_grade} is less than 4"
+
+    Rota.constraint_options["min_summed_grade_by_shifts_per_day"] = [
+        MinSummedGradeByShiftsPerDayConstraint(
+            shifts=["a", "b"],
+            min_grade_sum=5,
+        ),
+    ]
+
+    Rota.build_and_solve(options={"ratio": 0.000})
+    assert Rota.results.solver.status == "ok"
+
+    for worker in Rota.get_workers():
+        if worker.grade == 1:
+            assert Rota.get_worker_shift_list(worker).count("a") == 0, f"Worker {worker.name} with grade 1 should not have shift 'a'"
+            assert Rota.get_worker_shift_list(worker).count("b") == 0, f"Worker {worker.name} with grade 1 should not have shift 'b'"

test/test_limit_constraints.py

@@ -60,17 +60,17 @@ def test_constraint_limit_grades2(limit_constraint_rota):
             balance_offset=40,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "limit_grade_number", "options": {3: 2}}],
+            constraint=[{"name": "limit_grade_number", "options": {3: 1}}],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
     Rota.build_and_solve(options={"ratio": 0.1})
-    Rota.export_rota_to_html("test9")
+    Rota.export_rota_to_html("test9", folder="tests")
     grade_workers = Rota.get_workers_by_grade()
     for grade in grade_workers:
         shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
         zipped_lists = list(zip(*shift_patterns))
-        limit = 2 if grade == 3 else 2
+        limit = 1 if grade == 3 else 3
         for day_shifts in zipped_lists:
             assert day_shifts.count("weekday_night") == limit
     assert Rota.results.solver.status == "ok"

test/test_weekend_padding.py

@@ -46,7 +46,7 @@ class TestDemoRota:
 
         self.Rota.build_and_solve(options={"ratio": 0.00})
 
-        self.Rota.export_rota_to_html("balance_weekend")
+        self.Rota.export_rota_to_html("balance_weekend", folder="tests")
 
         assert self.Rota.results.solver.status == "ok"
         assert self.Rota.results.solver.termination_condition == "optimal"