a few more fixes and improvements

test/test_workers.py

@@ -1380,4 +1380,79 @@ def test_workers_double_shifts_on_weekends():
     Rota.build_and_solve(options={"ratio": 0.0})
     Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
 
-    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.status == "ok"
+
+
+def test_max_shifts_per_week_by_shift_name_basic():
+    Rota = RotaBuilder(weeks_to_rota=2, start_date=datetime.date(2025, 6, 16))
+    w = Worker(name="A01", site="site1", grade=1)
+    w.max_shifts_per_week_by_shift_name = {"a": 1}
+    Rota.add_worker(w)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("site1",),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    # Worker should not be assigned more than 1 "a" shift per week
+    for week in Rota.weeks:
+        count = sum(
+            Rota.model.works[w.id, week, day, "a"].value > 0.5
+            for day in days[:5]
+        )
+        assert count <= 1
+
+def test_max_shifts_per_week_by_shift_name_multiple_weeks():
+    Rota =generate_basic_rota(workers=4, weeks_to_rota=2)
+    for worker in Rota.get_workers():
+        worker.set_max_shifts_per_week("a", 1)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    assert Rota.results.solver.status == "error"
+
+    Rota.add_worker(
+        Worker(name="extraworker", site="group1", grade=1, max_shifts_per_week_by_shift_name={"a": 1})
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    assert Rota.results.solver.status == "ok"
+
+    Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
+    # Worker should not be assigned more than 2 "a" shifts per week
+    for worker in Rota.get_workers():
+        for week in Rota.weeks:
+            count = sum(
+                Rota.model.works[worker.id, week, day, "a"].value > 0.5
+                for day in days[:5]
+            )
+            assert count <= 2
+
+
+def test_max_shifts_per_week_by_shift_name_invalid_shift():
+    Rota = RotaBuilder(weeks_to_rota=1, start_date=datetime.date(2025, 6, 16))
+    w = Worker(name="A03", site="site1", grade=1)
+    # "b" does not exist
+    w.max_shifts_per_week_by_shift_name = {"b": 1}
+    Rota.add_worker(w)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("site1",),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=1,
+        ),
+    )
+    with pytest.raises(InvalidShift):
+        Rota.build_and_solve(options={"ratio": 0.0})