Ross
277 lines
12 KiB
Python
277 lines
12 KiB
Python
|
|
import pytest
|
|
import datetime
|
|
from rota_generator.shifts import RotaBuilder, SingleShift, days
|
|
from rota_generator.workers import Worker
|
|
|
|
@pytest.fixture
|
|
def basic_rota():
|
|
start_date = datetime.date(2022, 4, 4) # must be a Monday!
|
|
Rota = RotaBuilder(start_date=start_date, weeks_to_rota=2)
|
|
worker = Worker(name="worker1", site="site1", grade=1)
|
|
Rota.add_worker(worker)
|
|
return Rota
|
|
|
|
@pytest.fixture
|
|
def long_rota():
|
|
start_date = datetime.date(2022, 4, 4) # must be a Monday!
|
|
Rota = RotaBuilder(start_date=start_date, weeks_to_rota=10)
|
|
worker = Worker(name="worker1", site="site1", grade=1)
|
|
Rota.add_worker(worker)
|
|
return Rota
|
|
|
|
@pytest.fixture
|
|
def cons_rota():
|
|
start_date = datetime.date(2022, 4, 4) # must be a Monday!
|
|
Rota = RotaBuilder(start_date=start_date, weeks_to_rota=10)
|
|
worker1 = Worker(name="worker1", site="a")
|
|
worker2 = Worker(name="worker2", site="a")
|
|
worker3 = Worker(name="worker3", site="a")
|
|
worker4 = Worker(name="worker4", site="a")
|
|
worker5 = Worker(name="worker5", site="a")
|
|
worker6 = Worker(name="worker6", site="a")
|
|
worker7 = Worker(name="worker7", site="a")
|
|
worker8 = Worker(name="worker8", site="a")
|
|
worker9 = Worker(name="worker9", site="a")
|
|
worker10 = Worker(name="worker10", site="a")
|
|
worker11 = Worker(name="worker11", site="a")
|
|
|
|
worker12 = Worker(name="worker12", site="b")
|
|
worker13 = Worker(name="worker13", site="b")
|
|
worker14 = Worker(name="worker14", site="b")
|
|
worker15 = Worker(name="worker15", site="b")
|
|
worker16 = Worker(name="worker16", site="b")
|
|
worker17 = Worker(name="worker17", site="b")
|
|
worker18 = Worker(name="worker18", site="b")
|
|
worker19 = Worker(name="worker19", site="b")
|
|
worker20 = Worker(name="worker20", site="b")
|
|
|
|
Rota.add_workers([
|
|
worker1, worker2, worker3, worker4, worker5,
|
|
worker6, worker7, worker8, worker9, worker10,
|
|
worker11, worker12, worker13, worker14, worker15,
|
|
worker16, worker17, worker18, worker19, worker20
|
|
])
|
|
|
|
|
|
return Rota
|
|
|
|
def test_single_shift_per_day_default(basic_rota):
|
|
Rota = basic_rota
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["site1"], name="a", length=8, days=days[:5], workers_required=1),
|
|
SingleShift(sites=["site1"], name="b", length=8, days=days[:5], workers_required=1),
|
|
)
|
|
Rota.build_and_solve()
|
|
assert Rota.results.solver.status == "error"
|
|
# Only one shift per day should be assigned
|
|
|
|
def test_allow_two_shifts_per_day(basic_rota):
|
|
Rota = basic_rota
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["site1"], name="a", length=8, days=days[:2], workers_required=1),
|
|
SingleShift(sites=["site1"], name="b", length=8, days=days[:2], workers_required=1),
|
|
)
|
|
for worker in Rota.workers:
|
|
worker.allow_shifts_together("a", "b")
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_allow_two_shifts", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
|
|
# Now both shifts can be assigned together
|
|
for week, day in Rota.get_week_day_combinations():
|
|
assigned = [
|
|
shift for shift in Rota.get_shift_names_by_week_day(week, day)
|
|
if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5
|
|
]
|
|
assert len(assigned) <= 2
|
|
|
|
def test_only_specified_shifts_can_be_double_assigned(basic_rota):
|
|
Rota = basic_rota
|
|
worker2 = Worker(name="worker2", site="site1", grade=1)
|
|
Rota.add_worker(worker2)
|
|
#Rota.add_worker(worker3)
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["site1"], name="a", length=8, days=days[:3], workers_required=1),
|
|
SingleShift(sites=["site1"], name="b", length=8, days=days[:3], workers_required=1),
|
|
SingleShift(sites=["site1"], name="c", length=8, days=days[:3], workers_required=1),
|
|
)
|
|
for worker in Rota.workers:
|
|
worker.allow_shifts_together("a", "b")
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_only_specified_shifts_can_be_double_assigned", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
# Only "a" and "b" can be assigned together, never "a"+"c" or "b"+"c"
|
|
for worker in Rota.workers:
|
|
for week, day in Rota.get_week_day_combinations():
|
|
assigned = set([
|
|
shift for shift in Rota.get_shift_names_by_week_day(week, day)
|
|
if Rota.model.works[worker.id, week, day, shift].value > 0.8
|
|
])
|
|
|
|
assert assigned != set(["a", "c"]) and assigned != set(["b", "c"])
|
|
|
|
def test_triple_shifts_allowed(basic_rota):
|
|
Rota = basic_rota
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["site1"], name="a", length=8, days=days[2:4], workers_required=1),
|
|
SingleShift(sites=["site1"], name="b", length=8, days=days[2:4], workers_required=1),
|
|
SingleShift(sites=["site1"], name="c", length=8, days=days[2:4], workers_required=1),
|
|
)
|
|
Rota.allow_shifts_together_for_all_workers("a", "b", "c")
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_triple_shifts_allowed", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
def test_triple_shifts_allowed2(long_rota):
|
|
Rota = long_rota
|
|
worker2 = Worker(name="worker2", site="site1", grade=1)
|
|
worker3 = Worker(name="worker3", site="site1", grade=1)
|
|
Rota.add_worker(worker2)
|
|
Rota.add_worker(worker3)
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["site1"], name="a", length=8, days=days[2:4], workers_required=1, balance_offset=1),
|
|
SingleShift(sites=["site1"], name="b", length=8, days=days[2:4], workers_required=1, balance_offset=1),
|
|
SingleShift(sites=["site1"], name="c", length=8, days=days[2:4], workers_required=1, balance_offset=1),
|
|
)
|
|
|
|
worker2.allow_shifts_together("a", "b", "c")
|
|
worker2.prefer_multi_shift_together = 10
|
|
worker3.allow_shifts_together("a", "b", "c")
|
|
worker3.prefer_multi_shift_together = -10
|
|
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_triple_shifts_allowed2", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
def test_multi_shift_prefs(long_rota):
|
|
Rota = long_rota
|
|
#Rota.constraint_options["max_shifts_per_week"] = 10
|
|
worker2 = Worker(name="worker2", site="site1", grade=1)
|
|
worker3 = Worker(name="worker3", site="site1", grade=1)
|
|
worker4 = Worker(name="worker4", site="site1", grade=1)
|
|
#worker5 = Worker(name="worker5", site="site1", grade=1)
|
|
Rota.add_worker(worker2)
|
|
Rota.add_worker(worker3)
|
|
Rota.add_worker(worker4)
|
|
#Rota.add_worker(worker5)
|
|
Rota.add_worker
|
|
for worker in Rota.workers:
|
|
worker.allow_shifts_together("a", "b")
|
|
worker.prefer_multi_shift_together = 0
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["site1"], name="a", length=8, days=days[:1], workers_required=1),
|
|
SingleShift(sites=["site1"], name="b", length=8, days=days[:1], workers_required=1),
|
|
SingleShift(sites=["site1"], name="c", length=8, days=days[:1], workers_required=1),
|
|
SingleShift(sites=["site1"], name="d", length=8, days=days[:1], workers_required=1),
|
|
SingleShift(sites=["site1"], name="e", length=8, days=days[:1], workers_required=1),
|
|
#SingleShift(sites=["site1"], name="d", length=8, days=days[:1], workers_required=1),
|
|
)
|
|
Rota.build_and_solve(options={"ratio": 0})
|
|
Rota.export_rota_to_html("test_multi_shift_prefs", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
def test_cons_rota(cons_rota):
|
|
Rota = cons_rota
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["a"], name="oncall", length=8, days=days[:5], workers_required=1),
|
|
SingleShift(sites=["a", "b"], name="evening", length=8, days=days[:5], workers_required=1),
|
|
)
|
|
print("Adding shifts")
|
|
for worker in Rota.get_workers_by_site("a"):
|
|
print(worker.name)
|
|
worker.allow_shifts_together("oncall", "evening")
|
|
worker.prefer_multi_shift_together = 1
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_cons_rota", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
# Check that everyone at site "a" is assigned 2 or 3 evening shifts
|
|
for worker in Rota.get_workers_by_site("a"):
|
|
evening_count = Rota.get_worker_shift_count(worker, "evening")
|
|
assert 2 <= evening_count <= 3, f"{worker.name} assigned {evening_count} evening shifts"
|
|
|
|
if worker.site == "a":
|
|
oncall_count = Rota.get_worker_shift_count(worker, "oncall")
|
|
assert oncall_count == 4 or oncall_count == 5, f"{worker.name} assigned {oncall_count} oncall shifts"
|
|
|
|
# Check that days worked is no more than oncall count
|
|
days_worked = Rota.get_worker_days_worked(worker)
|
|
assert days_worked <= oncall_count, f"{worker.name} worked {days_worked} days, more than oncall shifts {oncall_count}"
|
|
|
|
def test_cons_rota2(cons_rota):
|
|
Rota = cons_rota
|
|
Rota.add_shifts(
|
|
SingleShift(sites=["a"], name="oncall", length=8, days=days[:5], workers_required=1),
|
|
SingleShift(sites=["a", "b"], name="evening", length=8, days=days[:5], workers_required=1),
|
|
)
|
|
print("Adding shifts")
|
|
for worker in Rota.get_workers_by_site("a"):
|
|
print(worker.name)
|
|
worker.allow_shifts_together("oncall", "evening")
|
|
worker.prefer_multi_shift_together = -1
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_cons_rota", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
# Check that everyone at site "a" is assigned 2 or 3 evening shifts
|
|
for worker in Rota.get_workers_by_site("a"):
|
|
evening_count = Rota.get_worker_shift_count(worker, "evening")
|
|
assert 2 <= evening_count <= 3, f"{worker.name} assigned {evening_count} evening shifts"
|
|
|
|
if worker.site == "a":
|
|
oncall_count = Rota.get_worker_shift_count(worker, "oncall")
|
|
assert oncall_count == 4 or oncall_count == 5, f"{worker.name} assigned {oncall_count} oncall shifts"
|
|
|
|
# Check that days worked is now more than oncall count
|
|
days_worked = Rota.get_worker_days_worked(worker)
|
|
assert days_worked > oncall_count, f"{worker.name} worked {days_worked} days, more than oncall shifts {oncall_count}"
|
|
|
|
def test_cons_rota3(cons_rota):
|
|
Rota = cons_rota
|
|
|
|
Rota.terminate_on_warning.remove("Worker/no valid shifts")
|
|
Rota.add_shifts(
|
|
#SingleShift(sites=["a"], name="oncall weekday", length=8, days=days[:4], workers_required=1,
|
|
# constraint=[
|
|
# {"name": "pre", "options": 2},
|
|
# {"name": "post", "options": 2},
|
|
# ],
|
|
# ),
|
|
##SingleShift(sites=["a"], name="oncall weekend", length=8, days=days[4:], workers_required=1),
|
|
#SingleShift(sites=["a", "b"], name="evening", length=8, days=days[:4], workers_required=1,
|
|
# constraint=[
|
|
# {"name": "pre", "options": 2},
|
|
# {"name": "post", "options": 2},
|
|
# ],
|
|
# ),
|
|
SingleShift(sites=["a"], name="help", length=8, days=days[4:], workers_required=1,
|
|
#assign_as_block=True,
|
|
#force_as_block=True
|
|
#constraint=[
|
|
# {"name": "pre", "options": 2},
|
|
# {"name": "post", "options": 2},
|
|
#],
|
|
) ,
|
|
SingleShift(sites=["b"], name="weekend b", length=8, days=days[5:], workers_required=1,
|
|
#constraint=[
|
|
# {"name": "pre", "options": 2},
|
|
# {"name": "post", "options": 2},
|
|
#],
|
|
),
|
|
)
|
|
#print("Adding shifts")
|
|
#for worker in Rota.get_workers_by_site("a"):
|
|
# print(worker.name)
|
|
## worker.allow_shifts_together("oncall weekday", "evening")
|
|
## #worker.allow_shifts_together("oncall weekend", "weekend")
|
|
## worker.prefer_multi_shift_together = 1
|
|
|
|
# worker.assign_as_block_preferences = {"help":10}
|
|
#for worker in Rota.get_workers_by_site("b"):
|
|
# worker.assign_as_block_preferences = {"weekend b":-10}
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_cons_rota", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|