Ross
95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
import datetime
|
|
import pytest
|
|
from rota_generator.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days, MinSummedGradeByShiftsPerDayConstraint
|
|
|
|
from rota_generator.workers import Worker
|
|
|
|
import itertools
|
|
|
|
|
|
def generate_basic_rota(weeks_to_rota=9):
|
|
start_date = datetime.date(2022, 3, 7)
|
|
|
|
Rota = RotaBuilder(
|
|
start_date,
|
|
weeks_to_rota=weeks_to_rota,
|
|
)
|
|
|
|
# Add a few workers
|
|
Rota.add_workers(
|
|
[
|
|
Worker(name="worker1", site="group1", grade=1),
|
|
Worker(name="worker2", site="group1", grade=2),
|
|
Worker(name="worker3", site="group1", grade=3),
|
|
Worker(name="worker4", site="group1", grade=1),
|
|
Worker(name="worker5", site="group1", grade=2),
|
|
Worker(name="worker6", site="group1", grade=3),
|
|
]
|
|
)
|
|
|
|
return Rota
|
|
|
|
|
|
def test_summed_shift_grade_constraints():
|
|
Rota = generate_basic_rota(5)
|
|
|
|
Rota.add_shifts(
|
|
SingleShift(
|
|
sites=("group1",), name="a", length=12.5, days=days,
|
|
workers_required=1,
|
|
force_as_block=False,
|
|
),
|
|
SingleShift(
|
|
sites=("group1",), name="b", length=12.5, days=days,
|
|
workers_required=1,
|
|
force_as_block=False,
|
|
),
|
|
)
|
|
|
|
# Add summed shift grade constraints
|
|
Rota.constraint_options["min_summed_grade_by_shifts_per_day"].append(
|
|
MinSummedGradeByShiftsPerDayConstraint(
|
|
shifts=["a", "b"],
|
|
min_grade_sum=4,
|
|
),
|
|
)
|
|
|
|
Rota.build_and_solve(options={"ratio": 0.000})
|
|
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
Rota.export_rota_to_html("summed_shift_grade_constraints", folder="tests")
|
|
|
|
a_grades = []
|
|
b_grades = []
|
|
|
|
for shift, shift_timetable in Rota.get_shift_timetable_by_week().items():
|
|
for week, week_timetable in shift_timetable.items():
|
|
for day, workers in week_timetable.items():
|
|
if shift == "a":
|
|
for worker in workers:
|
|
a_grades.append(worker.grade)
|
|
elif shift == "b":
|
|
for worker in workers:
|
|
b_grades.append(worker.grade)
|
|
|
|
# Check that the minimum summed grade is respected
|
|
summed_grades = [a + b for a, b in zip(a_grades, b_grades)]
|
|
for summed_grade in summed_grades:
|
|
assert summed_grade >= 4, f"Summed grade {summed_grade} is less than 4"
|
|
|
|
Rota.constraint_options["min_summed_grade_by_shifts_per_day"] = [
|
|
MinSummedGradeByShiftsPerDayConstraint(
|
|
shifts=["a", "b"],
|
|
min_grade_sum=5,
|
|
),
|
|
]
|
|
|
|
Rota.build_and_solve(options={"ratio": 0.000})
|
|
assert Rota.results.solver.status == "ok"
|
|
|
|
for worker in Rota.get_workers():
|
|
if worker.grade == 1:
|
|
assert Rota.get_worker_shift_list(worker).count("a") == 0, f"Worker {worker.name} with grade 1 should not have shift 'a'"
|
|
assert Rota.get_worker_shift_list(worker).count("b") == 0, f"Worker {worker.name} with grade 1 should not have shift 'b'"
|