import datetime import pytest from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days from rota.workers import Worker def generate_basic_rota(weeks_to_rota=10): start_date = datetime.date(2022, 3, 7) Rota = RotaBuilder( start_date, weeks_to_rota=weeks_to_rota, ) Rota.constraint_options["balance_weekends"] = True # Add a few workers Rota.add_workers( [ Worker(name="worker1", site="group1", grade=1), Worker(name="worker2", site="group1", grade=1), Worker(name="worker3", site="group2", grade=1), # Worker(name="worker4", site="group2", grade=1), # Worker(name="worker5", site="group2", grade=1), # Worker(name="worker6", site="group2", grade=1, fte=50), ] ) return Rota class TestBalancing: def test_basic_balancing(self): Rota = generate_basic_rota() Rota.add_shifts( SingleShift( sites=("group1", "group2"), name="a", length=12.5, days=days, balance_offset=2, workers_required=1, force_as_block=False, ), SingleShift( sites=("group1", "group2"), name="b", length=12.5, days=days, balance_offset=2, workers_required=1, force_as_block=False, ), ) Rota.build_and_solve(options={"ratio": 0.001}) Rota.export_rota_to_html("basic_balancing") 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) assert worker_shifts["a"] in (22, 23, 24, 25) assert worker_shifts["b"] in (22, 23, 24, 25) def test_weighted_shift_balancing(self): Rota = generate_basic_rota(20) Rota.add_shifts( SingleShift( sites=["group1", "group2"], name="a", length=12.5, days=days, balance_weighting=4, workers_required=1, force_as_block=False, constraint=[{"name": "preclear2"}, {"name": "postclear2"}], ), SingleShift( sites=("group1", "group2"), name="b", length=12.5, days=days, workers_required=1, force_as_block=False, ), SingleShift( sites=("group1", "group2"), name="c", length=12.5, days=days[0], workers_required=1, force_as_block=False, ), ) Rota.build_and_solve(options={"ratio": 0.001}) Rota.export_rota_to_html("basic_balancing_weighted_shifts") 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) assert worker_shifts["a"] in (46, 47, 48) # assert worker_shifts["b"] in (22, 23, 24, 25) def test_weighted_shift_balancing2(self): Rota = generate_basic_rota(23) Rota.add_shifts( SingleShift( sites=("group1", "group2"), name="a", length=12.5, days=days, # balance_weighting=4, workers_required=1, force_as_block=False, # constraint=["preclear2", "postclear2"], ), SingleShift( sites=("group1", "group2"), name="b", length=12.5, days=days, balance_weighting=4, workers_required=1, force_as_block=False, ), SingleShift( sites=("group1", "group2"), name="c", length=12.5, days=days[0], workers_required=1, force_as_block=False, ), ) Rota.build_and_solve(options={"ratio": 0.001}) Rota.export_rota_to_html("basic_balancing_weighted_shifts2") 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) # assert worker_shifts["a"] in (46,47,48) # assert worker_shifts["c"] in (46,47,48) assert worker_shifts["b"] in (52, 53, 54, 55) def test_weighted_shift_balancing3(self): Rota = generate_basic_rota(23) Rota.add_shifts( SingleShift( sites=("group1", "group2"), name="a", length=12.5, days=days, # balance_weighting=4, workers_required=1, force_as_block=False, # constraint=["preclear2", "postclear2"], ), SingleShift( sites=("group1", "group2"), name="b", length=12.5, days=days, balance_weighting=4, workers_required=1, force_as_block=False, ), SingleShift( sites=("group1", "group2"), name="c", length=12.5, days=days[0], balance_weighting=4, workers_required=1, force_as_block=False, ), ) Rota.build_and_solve(options={"ratio": 0.000}) Rota.export_rota_to_html("basic_balancing_weighted_shifts3") 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) # assert worker_shifts["a"] in (46,47,48) assert worker_shifts["c"] in (7, 8) assert worker_shifts["b"] in (52, 53, 54, 55) def test_weighted_shift_balancing4(self): Rota = generate_basic_rota(10) Rota.add_worker( Worker(name="worker4", site="group3", grade=1), ) Rota.add_shifts( SingleShift( sites=("group1", "group2"), name="a", length=12.5, days=days, balance_offset=99, balance_weighting=8, workers_required=2, force_as_block=False, # constraint=["preclear2", "postclear2"], ), SingleShift( sites=("group2", "group3"), name="b", length=12.5, days=days, balance_offset=99, # balance_weighting=4, workers_required=1, force_as_block=False, ), ) Rota.build_and_solve(options={"ratio": 0.000}) Rota.export_rota_to_html("basic_balancing_weighted_shifts4") 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) assert worker_shifts["a"] in (0, 46, 47, 48) def test_weighted_shift_balancing5(self): Rota = generate_basic_rota(10) Rota.add_worker( Worker(name="worker4", site="group3", grade=1), ) Rota.add_shifts( SingleShift( sites=("group1", "group2"), name="a", length=12.5, days=days, balance_offset=99, # balance_weighting=8, workers_required=2, force_as_block=False, # constraint=["preclear2", "postclear2"], ), SingleShift( sites=("group2", "group3"), name="b", length=12.5, days=days, balance_offset=99, balance_weighting=4, workers_required=1, force_as_block=False, ), ) Rota.build_and_solve(options={"ratio": 0.000}) Rota.export_rota_to_html("basic_balancing_weighted_shifts5") 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) assert worker_shifts["b"] in (0, 35)