Ross
106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
import datetime
|
|
import pytest
|
|
from pytest import approx
|
|
from rota_generator.shifts import PostShiftConstraint, RotaBuilder, SingleShift, days
|
|
from rota_generator.workers import Worker
|
|
from rota_generator.shifts import PreShiftConstraint
|
|
|
|
@pytest.fixture
|
|
def demo_rota_nights():
|
|
weeks_to_rota = 12
|
|
start_date = datetime.date(2022, 3, 14)
|
|
|
|
Rota = RotaBuilder(
|
|
start_date,
|
|
weeks_to_rota=weeks_to_rota,
|
|
balance_offset_modifier=1,
|
|
)
|
|
|
|
Rota.constraint_options_model.max_shifts_per_week = 5
|
|
Rota.constraint_options_model.max_shifts_per_month = 31
|
|
# legacy option `limit_to_1_st2_on_nights` removed; test omits it
|
|
Rota.constraint_options_model.constrain_time_off_after_nights = False
|
|
Rota.constraint_options_model.balance_nights_across_sites = False
|
|
Rota.constraint_options_model.balance_shifts = True
|
|
Rota.constraint_options_model.balance_shifts_over_workers = True
|
|
Rota.constraint_options_model.balance_nights = False
|
|
Rota.constraint_options_model.minimise_shift_diffs = False
|
|
Rota.constraint_options_model.balance_blocks = False
|
|
Rota.constraint_options_model.max_night_frequency = 0
|
|
|
|
# Add workers
|
|
workers = [
|
|
Worker(name="worker1", site="group1", grade=1),
|
|
Worker(name="worker2", site="group1", grade=1),
|
|
Worker(name="worker3", site="group1", grade=1),
|
|
Worker(name="worker4", site="group1", grade=1),
|
|
Worker(name="worker5", site="group2", grade=1, fte=60),
|
|
Worker(name="worker6", site="group2", grade=1, fte=40),
|
|
Worker(name="worker7", site="group2", grade=1),
|
|
Worker(name="worker8", site="group2", grade=1),
|
|
]
|
|
Rota.add_workers(workers)
|
|
|
|
# Add shifts
|
|
Rota.add_shifts(
|
|
SingleShift(
|
|
sites=("group1", "group2"),
|
|
name="night_weekday",
|
|
length=12.5,
|
|
days=days[:4],
|
|
workers_required=1,
|
|
constraints=[
|
|
PreShiftConstraint(days=1),
|
|
PostShiftConstraint(days=2),
|
|
],
|
|
force_as_block=True
|
|
),
|
|
SingleShift(
|
|
sites=("group1", "group2"),
|
|
name="night_weekend",
|
|
length=12.5,
|
|
days=days[4:],
|
|
constraints=[
|
|
PreShiftConstraint(days=1),
|
|
PostShiftConstraint(days=2),
|
|
],
|
|
force_as_block=True,
|
|
),
|
|
SingleShift(
|
|
sites=("group1", "group2"),
|
|
name="twilight",
|
|
length=12.5,
|
|
days=days[:5],
|
|
workers_required=2,
|
|
),
|
|
)
|
|
|
|
Rota.build_shifts()
|
|
Rota.build_workers()
|
|
Rota.build_model()
|
|
|
|
solver_options = {"ratio": 0.01, "seconds": 1000, "threads": 10}
|
|
Rota.solve_model(options=solver_options)
|
|
Rota.export_rota_to_html("test3")
|
|
|
|
return Rota
|
|
|
|
def test_night_assignment(demo_rota_nights):
|
|
Rota = demo_rota_nights
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
shift_summary = Rota.get_shift_summary_dict()
|
|
|
|
for worker_name in shift_summary:
|
|
worker_shifts = shift_summary[worker_name]
|
|
worker = Rota.get_worker_by_name(worker_name)
|
|
|
|
if worker.fte == 40:
|
|
assert worker_shifts["night_weekday"] in (approx(0), approx(4))
|
|
assert worker_shifts["night_weekend"] in (approx(0), approx(3))
|
|
elif worker.fte == 60:
|
|
assert worker_shifts["night_weekday"] in (approx(0), approx(4), approx(8))
|
|
assert worker_shifts["night_weekend"] in (approx(0), approx(3), approx(6))
|
|
else:
|
|
assert worker_shifts["night_weekday"] in (approx(4), approx(8))
|
|
assert worker_shifts["night_weekend"] in (approx(3), approx(6)) |