Ross
e6bdd4a09e
- 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.
121 lines
5.1 KiB
Python
121 lines
5.1 KiB
Python
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")
|
|
|