Files
proc-rota/test/test_shift_constraints.py
T
2025-12-09 14:29:54 +00:00

71 lines
2.6 KiB
Python

import pytest
import datetime
from copy import deepcopy
from rota_generator.shifts import RotaBuilder, SingleShift, days
from rota_generator.workers import Worker
@pytest.fixture
def rota_2workers():
weeks_to_rota = 10
start_date = datetime.date(2022, 3, 7)
Rota = RotaBuilder(
start_date,
weeks_to_rota=weeks_to_rota,
balance_offset_modifier=1,
)
worker1 = Worker(name="worker1", site="group1", grade=1)
worker2 = Worker(name="worker2", site="group1", grade=1)
Rota.add_workers((worker1, worker2))
Rota.add_shifts(
SingleShift(
sites=("group1",),
name="a",
length=12.5,
days=days,
balance_offset=40,
workers_required=1,
),
)
return Rota
@pytest.fixture
def rota_3workers(rota_2workers):
Rota3 = deepcopy(rota_2workers)
worker3 = Worker(name="worker3", site="group1", grade=1)
Rota3.add_worker(worker3)
return Rota3
def test_max_shifts(rota_2workers):
rota_2workers.constraint_options_model.max_shifts_per_month = 14
rota_2workers.build_and_solve()
assert rota_2workers.results.solver.status == "ok"
assert rota_2workers.results.solver.termination_condition == "optimal"
def test_max_shifts_fail(rota_2workers):
rota_2workers.constraint_options_model.max_shifts_per_month = 13
rota_2workers.build_and_solve()
assert rota_2workers.results.solver.status in ("warning", "error")
assert rota_2workers.results.solver.termination_condition == "infeasible"
def test_max_shifts_extra_worker(rota_3workers):
rota_3workers.constraint_options_model.max_shifts_per_month = 12
rota_3workers.build_and_solve()
assert rota_3workers.results.solver.status == "ok"
assert rota_3workers.results.solver.termination_condition == "optimal"
def test_max_shifts_per_week_fail(rota_2workers):
rota_2workers.constraint_options_model.max_shifts_per_week = 3
rota_2workers.build_and_solve()
assert rota_2workers.results.solver.status in ("warning", "error")
assert rota_2workers.results.solver.termination_condition == "infeasible"
def test_max_shifts_per_week_pass(rota_2workers):
rota_2workers.constraint_options_model.max_shifts_per_month = 14
rota_2workers.constraint_options_model.max_shifts_per_week = 4
rota_2workers.build_and_solve()
assert rota_2workers.results.solver.status == "ok"
def test_max_shifts_per_week_extra_worker_pass(rota_3workers):
rota_3workers.constraint_options_model.max_shifts_per_week = 3
rota_3workers.build_and_solve()
assert rota_3workers.results.solver.status == "ok"