Add support for avoiding shifts on specific dates for workers

- Introduced AvoidShiftOnDates model to define constraints for workers. - Enhanced Worker class to register avoid-shift rules with the rota builder. - Implemented logic in RotaBuilder to enforce these constraints during shift assignment. - Added tests to verify functionality for avoiding shifts on single dates and date ranges.
2026-01-05 21:44:34 +00:00
parent 62e48ef205
commit e6bdd4a09e
6 changed files with 289 additions and 3 deletions
@@ -0,0 +1,120 @@
+import datetime
+from rota_generator.shifts import RotaBuilder, SingleShift, days
+from rota_generator.workers import Worker, AvoidShiftOnDates
+
+
+def build_simple_rota():
+    start = datetime.date(2026, 4, 6)  # Monday
+    Rota = RotaBuilder(start, weeks_to_rota=4, name="test_rota")
+    Rota.add_shifts(
+        SingleShift(sites=("exeter",), name="night_weekday", length=12.25, days=days[:4], workers_required=1),
+        SingleShift(sites=("exeter",), name="plymouth_twilight", length=8, days=days[5:], workers_required=1),
+    )
+    Rota.build_shifts()
+    return Rota
+
+
+def test_avoid_single_date_for_worker():
+    Rota = build_simple_rota()
+    # create a worker who should avoid night_weekday on 2026-04-07
+    avoid = AvoidShiftOnDates(shifts=["night_weekday"], dates=[datetime.date(2026, 4, 7)])
+    w = Worker(name="Test Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
+    Rota.add_worker(w)
+    Rota.build_workers()
+
+    # find week/day corresponding to 2026-04-07
+    weeks_days = [(wk, d) for wk, d in Rota.weeks_days_product if Rota.week_day_date_map[(wk, d)] == datetime.date(2026, 4, 7)]
+    assert len(weeks_days) == 1
+    wk, day = weeks_days[0]
+
+    # The model should mark works[worker.id, wk, day, 'night_weekday'] impossible -> check unavailable_to_work
+    assert (w.id, wk, day) not in Rota.unavailable_to_work
+
+    # Build and solve minimal model (but we can inspect constraints). Instead, check that the avoid rule was registered
+    found = False
+    for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates:
+        if "names" in c and w.name in c["names"]:
+            found = True
+            # ensure the date is represented
+            assert c.get("dates") is not None and datetime.date(2026,4,7) in c.get("dates")
+    assert found
+
+
+def test_avoid_date_range_for_worker_applies_to_multiple_days():
+    Rota = build_simple_rota()
+    # Avoid night_weekday for 2026-04-06 through 2026-04-09
+    avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="09/04/2026")
+    w = Worker(name="Range Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
+    Rota.add_worker(w)
+    Rota.build_workers()
+
+    # Collect dates between start and end
+    sd = datetime.date(2026,4,6)
+    ed = datetime.date(2026,4,9)
+    dates_in_range = [d for d in Rota.get_date_range(sd, ed)]
+    assert len(dates_in_range) == 4
+
+    # Ensure the rota-level constraint was added
+    found = False
+    for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates:
+        if w.name in c.get("names", []):
+            found = True
+            assert c.get("start_date") == sd
+            assert c.get("end_date") == ed
+    assert found
+
+
+def test_export_rota_for_visual_inspection():
+    Rota = build_simple_rota()
+    # add two workers
+    avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="26/04/2026")
+    w1 = Worker(name="Export A", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
+    w2 = Worker(name="Export B", site="exeter", grade=3, fte=100)
+
+
+
+    Rota.add_workers((w1, w2))
+
+    # Build model but do not solve, then export
+    Rota.name = "export_test_rota"
+    Rota.build_and_solve(solve=True)
+
+    assert Rota.results.solver.status == "ok", "Rota should be feasible"
+
+    # Next we check that worker w1 has no night_weekday shifts assigned between 06/04/2026 and 26/04/2026
+    sd = datetime.date(2026,4,6)
+    ed = datetime.date(2026,4,26)
+    for single_date in Rota.get_date_range(sd, ed):
+        wk_day = Rota.get_week_day_by_date(single_date)
+        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk_day[0], wk_day[1])
+        for sh in assigned_shifts:
+            assert sh.name != "night_weekday", f"Worker {w1.name} was assigned forbidden shift on {single_date}"
+
+    # We also check that w1 has 4 night shifts in total (the max possible in 4 weeks)
+    night_shift_count = 0
+    for wk, day in Rota.weeks_days_product:
+        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk, day)
+        for sh in assigned_shifts:
+            if sh.name == "night_weekday":
+                night_shift_count += 1
+    assert night_shift_count == 4, f"Worker {w1.name} should have 4 night shifts, has {night_shift_count}"
+
+    # And that w2 has 12 night shifts (no restrictions)
+    night_shift_count_w2 = 0
+    for wk, day in Rota.weeks_days_product:
+        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w2.id, wk, day)
+        for sh in assigned_shifts:
+            if sh.name == "night_weekday":
+                night_shift_count_w2 += 1
+    assert night_shift_count_w2 == 12, f"Worker {w2.name} should have 12 night shifts, has {night_shift_count_w2}"
+
+    # Next check that adding a short avoid to w2 results in an infeasible model
+    avoid2 = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="07/04/2026")
+    w2.avoid_shifts_on_dates.append(avoid2)
+    Rota.build_and_solve(solve=True)
+
+    assert Rota.results.solver.termination_condition == "infeasible", "Rota should be infeasible due to avoid shifts on dates constraints"
+
+
+    Rota.export_rota_to_html("avoid_shifts_on_dates_test.html", folder="tests")
+