Ross
87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
|
|
import pytest
|
|
import datetime
|
|
from rota.shifts import RotaBuilder, SingleShift, days
|
|
from rota.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
|
|
|
|
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
|
|
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) <= 1
|
|
|
|
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),
|
|
)
|
|
Rota.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),
|
|
)
|
|
Rota.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("a", "b", "c")
|
|
Rota.build_and_solve()
|
|
Rota.export_rota_to_html("test_triple_shifts_allowed", folder="tests")
|
|
assert Rota.results.solver.status == "ok"
|
|
|