ross/proc-rota
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
7fe44df476b79becc3be4dd00eb8b3a8a080ebdd
proc-rota/test/test_shifts.py
T
Ross 76f3b48c1d Refactor shift constraints in tests to use new constraint classes 
- Updated test cases in `test_workers.py` to replace dictionary-based constraints with instances of `PreShiftConstraint`, `PostShiftConstraint`, and `MaxShiftsPerWeekConstraint`.
- Added a new test `test_worker_assign_as_block_preference2` to verify worker preferences for block assignments.
- Marked tests as slow where appropriate to improve test suite performance.
2025-08-11 13:03:08 +01:00

455 lines
22 KiB
Python
Raw Blame History

import datetime
import pytest
from rota.shifts import (
InvalidShift, MaxShiftsPerWeekBlockConstraint, MaxShiftsPerWeekConstraint,
NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift,
WarningTermination, days, WorkerRequirement
)
from rota.workers import NotAvailableToWork, WorkRequests, Worker, generate_not_available_to_works
def generate_basic_rota(weeks_to_rota=10, workers=2, start_date=datetime.date(2022, 3, 7)):
Rota = RotaBuilder(
start_date,
weeks_to_rota=weeks_to_rota,
)
Rota.add_workers(
[Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)]
)
return Rota
def date_generator(from_date, days):
n = 0
while True:
yield from_date
n += 1
if n >= days:
break
from_date = from_date + datetime.timedelta(days=1)
def test_duplicate_shifts():
Rota = generate_basic_rota()
Rota.add_worker(Worker(name="worker3", site="group1", grade=1))
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days))
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
Rota.add_shifts(SingleShift(sites=("group1",), name="a", length=12.5, days=days))
with pytest.raises(InvalidShift):
Rota.build_and_solve(options={"ratio": 0.000})
def test_max_shifts():
Rota = generate_basic_rota()
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]))
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
def test_max_shifts_fail():
Rota = generate_basic_rota()
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=3)]))
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.termination_condition == "infeasible"
Rota.add_worker(Worker(name="worker3", site="group1", grade=1))
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
def test_max_shifts_per_week_block():
Rota = generate_basic_rota()
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2)]))
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("max_shifts_per_week_block")
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
def test_max_shifts_per_week_block_fail():
Rota = generate_basic_rota()
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekBlockConstraint(week_block=3)]))
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("max_shifts_per_week_block")
assert Rota.results.solver.status in ("warning", "error")
def test_max_shifts_per_week_block_shift_number():
Rota = generate_basic_rota()
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2, max_shifts=7)]))
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("max_shifts_per_week_block")
assert Rota.results.solver.status == "ok"
def test_max_shifts_per_week_block_shift_number_fail():
Rota = generate_basic_rota()
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2, max_shifts=6)]))
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("max_shifts_per_week_block")
assert Rota.results.solver.status in ("warning", "error")
def test_pre_shift_constraint():
Rota = generate_basic_rota()
for i, d in enumerate(days[:6]):
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name=chr(97+i), length=12.5, days=d,
constraints=[PreShiftConstraint(days=1)]))
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("pre")
assert Rota.results.solver.status == "ok"
for worker in Rota.get_workers():
shift_string = Rota.get_worker_shift_list_string(worker)
for s in ["ab", "bc", "cd", "de", "ef", "fg"]:
assert s not in shift_string
Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="g", length=12.5, days=days[6],
constraints=[PreShiftConstraint(days=2)]))
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.status in ("warning", "error")
def test_pre_shift_constraint2():
Rota = generate_basic_rota()
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
constraints=[PreShiftConstraint(days=1)], workers_required=2)
)
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("pre")
assert Rota.results.solver.status == "ok"
def test_pre_shift_constraint3():
Rota = generate_basic_rota()
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
constraints=[PreShiftConstraint(days=3)], workers_required=2),
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
constraints=[PreShiftConstraint(days=3)], workers_required=2),
)
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("pre")
assert Rota.results.solver.status in ("warning", "error")
def test_post_shift_constraint():
Rota = generate_basic_rota()
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
constraints=[PostShiftConstraint(days=3)], workers_required=2),
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
constraints=[PostShiftConstraint(days=2)], workers_required=2),
)
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("post")
assert Rota.results.solver.status == "ok"
def test_post_shift_constraint2():
Rota = generate_basic_rota()
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
constraints=[PostShiftConstraint(days=3)], workers_required=2),
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
constraints=[PostShiftConstraint(days=3)], workers_required=2),
)
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("post")
assert Rota.results.solver.status in ("warning", "error")
def test_shift_constraint_week_limit():
Rota = generate_basic_rota()
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
constraints=[PostShiftConstraint(days=1, weeks=[1,2,3])], workers_required=2),
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[1], workers_required=2),
)
Rota.build_and_solve(options={"ratio": 0.001})
Rota.export_rota_to_html("post")
assert Rota.results.solver.status in ("warning", "error")
Rota.add_workers([
Worker(name="worker3", end_date=Rota.start_date + datetime.timedelta(weeks=3), site="group1", grade=1),
Worker(name="worker4", end_date=Rota.start_date + datetime.timedelta(weeks=3), site="group1", grade=1)
])
Rota.build_and_solve(options={"ratio": 0.001})
assert Rota.results.solver.status == "ok"
def test_shift_constraint_week_limit2():
Rota = generate_basic_rota()
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
constraints=[PostShiftConstraint(days=1, start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4))], workers_required=2),
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[1], workers_required=2),
)
Rota.build_and_solve(options={"ratio": 0.001})
assert Rota.results.solver.status in ("warning", "error")
Rota.add_workers([
Worker(name="worker3", start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4), site="group1", grade=1),
Worker(name="worker4", start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4), site="group1", grade=1)
])
Rota.build_and_solve(options={"ratio": 0.001})
Rota.export_rota_to_html("post")
assert Rota.results.solver.status == "ok"
def test_shift_start_date():
Rota = generate_basic_rota(workers=2)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
start_date=(Rota.start_date + datetime.timedelta(days=7))),
)
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
assert Rota.get_workers_total_shifts()["worker1"] in (31, 32)
assert Rota.get_workers_total_shifts()["worker2"] in (31, 32)
def test_shift_start_date_end_date():
Rota = generate_basic_rota(workers=2)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
start_date=(Rota.start_date + datetime.timedelta(days=7)),
end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
)
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
for worker in Rota.get_workers():
assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)
def test_shift_invalid_start_date():
Rota = generate_basic_rota(workers=2)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
start_date=(Rota.start_date - datetime.timedelta(days=6))),
)
with pytest.raises(WarningTermination):
Rota.build_and_solve(options={"ratio": 0.000})
def test_shift_invalid_start_date2():
Rota = generate_basic_rota(workers=2)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
start_date=(Rota.rota_end_date + datetime.timedelta(days=1))),
)
with pytest.raises(WarningTermination):
Rota.build_and_solve(options={"ratio": 0.000})
def test_shift_invalid_end_date():
Rota = generate_basic_rota(workers=2)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
end_date=(Rota.rota_end_date + datetime.timedelta(days=1))),
)
with pytest.raises(WarningTermination):
Rota.build_and_solve(options={"ratio": 0.000})
def test_shift_invalid_end_date2():
Rota = generate_basic_rota(workers=2)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
end_date=(Rota.start_date - datetime.timedelta(days=1))),
)
with pytest.raises(WarningTermination):
Rota.build_and_solve(options={"ratio": 0.000})
def test_shift_start_date_end_date_block():
Rota = generate_basic_rota(workers=2)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
assign_as_block=True,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
start_date=(Rota.start_date + datetime.timedelta(days=7)),
end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
)
Rota.build_and_solve(options={"ratio": 0.000})
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
for worker in Rota.get_workers():
assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)
def test_shift_start_date_end_date_block_more_workers():
Rota = generate_basic_rota(workers=10)
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
assign_as_block=True, force_as_block=True, workers_required=4,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]),
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
assign_as_block=True, force_as_block=True, workers_required=4,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
start_date=(Rota.start_date + datetime.timedelta(days=7)),
end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
)
Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("test_shift_start_date")
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
total_shifts = Rota.get_workers_total_shifts()
for worker in Rota.get_workers():
assert total_shifts[worker.name] == 22
shift_string = Rota.get_worker_shift_list_string(worker)
assert "ab" not in shift_string
assert "ba" not in shift_string
def test_shift_start_date_end_date_block_with_requests():
Rota = generate_basic_rota(workers=7)
request_date = Rota.start_date + datetime.timedelta(days=(7*6)+4)
request = [WorkRequests(date=request_date, shift="c")]
Rota.add_workers([
Worker(name=f"worker-wr-{i}", site="group2", grade=1, work_requests=request) for i in range(1, 4)
])
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
assign_as_block=True, force_as_block=True, workers_required=4,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]),
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
assign_as_block=True, force_as_block=True, workers_required=4,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
start_date=(Rota.start_date + datetime.timedelta(days=7)),
end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
SingleShift(sites=("group1", "group2"), name="c", length=12.5, days=days[4:],
assign_as_block=True, force_as_block=True, workers_required=4,
constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
start_date=(Rota.start_date + datetime.timedelta(days=(7*6)))),
)
Rota.build_and_solve(options={"ratio": 0.1})
Rota.export_rota_to_html("test_shift_start_date")
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
total_shifts = Rota.get_workers_total_shifts()
for worker in Rota.get_workers():
assert total_shifts[worker.name] in range(24, 30)
shift_string = Rota.get_worker_shift_list_string(worker)
assert "ab" not in shift_string
assert "ba" not in shift_string
assert shift_string.count("aaaa") == shift_string.count("a") / 4
assert shift_string.count("bbb") == shift_string.count("b") / 3
assert shift_string.count("ccc") == shift_string.count("c") / 3
for worker in Rota.get_workers_by_group()["group2"]:
assert Rota.get_worker_shifts_by_date(worker)[request_date] == "c"
def test_worker_requirement_start_and_end_date():
WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14))
def test_worker_requirement_end_date_before_start_date():
with pytest.raises(ValueError):
WorkerRequirement(start_date=datetime.date(2022, 3, 14), end_date=datetime.date(2022, 3, 7))
def test_worker_requirement():
Rota = generate_basic_rota(workers=2)
wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21))]
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
assign_as_block=True, force_as_block=True, workers_required=wr),
)
Rota.build_and_solve(options={"ratio": 0.1})
Rota.export_rota_to_html("test_worker_requirement")
assert Rota.results.solver.status == "ok"
total_shifts = Rota.get_workers_total_shifts()
for worker in Rota.get_workers():
assert total_shifts[worker.name] == 4
assert "aaaa" in Rota.get_worker_shift_list_string(worker)
def test_worker_requirement2():
Rota = generate_basic_rota(workers=2)
wr = [
WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21)),
WorkerRequirement(start_date=datetime.date(2022, 3, 28))
]
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
assign_as_block=True, workers_required=wr),
)
Rota.build_and_solve(options={"ratio": 0.1})
Rota.export_rota_to_html("test_worker_requirement")
assert Rota.results.solver.status == "ok"
total_shifts = Rota.get_workers_total_shifts()
for worker in Rota.get_workers():
assert total_shifts[worker.name] == 18
assert Rota.get_worker_shift_list_string(worker)[14:21] == "-------"
def test_worker_requirement_increase_workers():
Rota = generate_basic_rota(workers=2)
wr = [
WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
]
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
assign_as_block=True, workers_required=wr),
)
Rota.build_and_solve(options={"ratio": 0.1})
Rota.export_rota_to_html("test_worker_requirement")
assert Rota.results.solver.status == "ok"
total_shifts = Rota.get_workers_total_shifts()
for worker in Rota.get_workers():
assert total_shifts[worker.name] == 30
assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---"*7)
def test_worker_requirement_balancing():
Rota = generate_basic_rota(workers=2)
natws = generate_not_available_to_works(datetime.date(2022,4,18), days=21, reason="holiday")
Rota.add_worker(Worker(name="worker3", site="group1", grade=1, not_available_to_work=natws))
wr = [
WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
]
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
assign_as_block=True, workers_required=wr),
)
Rota.build_and_solve(options={"ratio": 0.0})
Rota.export_rota_to_html("test_worker_requirement")
assert Rota.results.solver.status == "ok"
total_shifts = Rota.get_workers_total_shifts()
for worker in Rota.get_workers():
assert total_shifts[worker.name] == 20
assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
if worker.name == "worker3":
assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---")
assert Rota.get_worker_shift_list_string(worker).startswith("aaaa---")
def test_worker_requirement_multiple_shifts():
Rota = generate_basic_rota(workers=4)
wr = [
WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
]
Rota.add_shifts(
SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
assign_as_block=True, workers_required=2),
SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
assign_as_block=True, workers_required=wr),
)
Rota.build_and_solve(options={"ratio": 0.1})
Rota.export_rota_to_html("test_worker_requirement")
assert Rota.results.solver.status == "ok"
total_shifts = Rota.get_workers_total_shifts()
for worker in Rota.get_workers():
assert total_shifts[worker.name] == 30
def test_display_char_in_output(monkeypatch):
Rota = generate_basic_rota(workers=1)
shift = SingleShift(
sites=("group1",),
name="testshift",
length=8,
days=days[:1],
display_char="Z"
)
Rota.add_shifts(shift)
assert shift.display_char == "Z"
Rota.build_and_solve(options={"ratio": 0.1})
assert Rota.results.solver.status == "ok"
for worker in Rota.get_workers():
assert Rota.get_worker_shift_list_string(worker).startswith("Z------" * 5)
Reference in New Issue View Git Blame Copy Permalink