Files
proc-rota/test/test_balance_shift_sites.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

139 lines
4.8 KiB
Python

import pytest
import datetime
from rota.shifts import BalanceAcrossGroupsConstraint, RotaBuilder, SingleShift, days
from rota.workers import Worker
@pytest.fixture
def demo_rota_balance_sites():
weeks_to_rota = 8
start_date = datetime.date(2022, 3, 7)
Rota = RotaBuilder(
start_date,
weeks_to_rota=weeks_to_rota,
)
worker1 = Worker(name="worker1", site="group1", grade=1)
worker2 = Worker(name="worker2", site="group1", grade=1)
worker3 = Worker(name="worker3", site="group2", grade=1)
worker4 = Worker(name="worker4", site="group2", grade=1)
worker5 = Worker(name="worker5", site="group3", grade=1)
worker6 = Worker(name="worker6", site="group3", grade=1)
Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6))
return Rota
def test_balance_blocks_across_2_groups(demo_rota_balance_sites):
Rota = demo_rota_balance_sites
Rota.add_shifts(
SingleShift(
sites=("group1", "group2"),
name="weekday_night",
length=12.5,
days=days[:4],
balance_offset=40,
workers_required=2,
force_as_block=True,
constraints=[BalanceAcrossGroupsConstraint()],
),
)
Rota.constraint_options["balance_nights_across_sites"] = False
Rota.constraint_options["balance_nights"] = False
Rota.constraint_options["constrain_time_off_after_nights"] = False
# Worker 5 and 6 don't have any valid shifts so this should fail
#with pytest.raises(WarningTermination):
# Rota.build_and_solve(options={"ratio": 0.0})
# Unless we remove the warning
Rota.terminate_on_warning.remove("Worker/no valid shifts")
Rota.build_and_solve(options={"ratio": 0.0})
assert len(Rota.get_warnings("Worker/no valid shifts")) == 2
Rota.export_rota_to_html("test5")
group_workers = Rota.get_workers_by_group()
for group in group_workers:
shift_patterns = []
for w in group_workers[group]:
shift_patterns.append(Rota.get_worker_shift_list(w))
zipped_lists = list(zip(*shift_patterns))
for day_shifts in zipped_lists:
assert day_shifts.count("weekday_night") <= 1
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
def test_balance_blocks_across_2_groups_unbalanced(demo_rota_balance_sites):
Rota = demo_rota_balance_sites
Rota.shifts = []
Rota.add_shifts(
SingleShift(
sites=("group1", "group2"),
name="weekday_night",
length=12.5,
days=days[:4],
balance_offset=40,
workers_required=3,
force_as_block=True,
constraints=[BalanceAcrossGroupsConstraint()],
),
)
Rota.constraint_options["balance_nights_across_sites"] = False
Rota.constraint_options["balance_nights"] = False
Rota.constraint_options["constrain_time_off_after_nights"] = False
try:
Rota.terminate_on_warning.remove("Worker/no valid shifts")
except ValueError:
# Seems to happen with sequential test running
pass
Rota.build_and_solve(options={"ratio": 0.1})
group_workers = Rota.get_workers_by_group()
for group in group_workers:
shift_patterns = []
for w in group_workers[group]:
shift_patterns.append(Rota.get_worker_shift_list(w))
zipped_lists = list(zip(*shift_patterns))
for day_shifts in zipped_lists:
assert day_shifts.count("weekday_night") <= 2
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"
def test_balance_blocks_across_3_groups(demo_rota_balance_sites):
Rota = demo_rota_balance_sites
Rota.shifts = []
Rota.add_shifts(
SingleShift(
sites=("group1", "group2", "group3"),
name="weekday_night",
length=12.5,
days=days[:4],
balance_offset=40,
workers_required=3,
force_as_block=True,
constraints=[BalanceAcrossGroupsConstraint()],
),
)
Rota.constraint_options["balance_nights_across_sites"] = False
Rota.constraint_options["balance_nights"] = False
Rota.constraint_options["constrain_time_off_after_nights"] = False
Rota.build_and_solve(options={"ratio": 0.1})
group_workers = Rota.get_workers_by_group()
for group in group_workers:
shift_patterns = []
for w in group_workers[group]:
shift_patterns.append(Rota.get_worker_shift_list(w))
zipped_lists = list(zip(*shift_patterns))
for day_shifts in zipped_lists:
assert day_shifts.count("weekday_night") <= 1
assert Rota.results.solver.status == "ok"
assert Rota.results.solver.termination_condition == "optimal"