import pytest import datetime from rota_generator.shifts import RotaBuilder, SingleShift, days from rota_generator.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 @pytest.fixture def long_rota(): start_date = datetime.date(2022, 4, 4) # must be a Monday! Rota = RotaBuilder(start_date=start_date, weeks_to_rota=10) worker = Worker(name="worker1", site="site1", grade=1) Rota.add_worker(worker) return Rota @pytest.fixture def cons_rota(): start_date = datetime.date(2022, 4, 4) # must be a Monday! Rota = RotaBuilder(start_date=start_date, weeks_to_rota=10) worker1 = Worker(name="worker1", site="a") worker2 = Worker(name="worker2", site="a") worker3 = Worker(name="worker3", site="a") worker4 = Worker(name="worker4", site="a") worker5 = Worker(name="worker5", site="a") worker6 = Worker(name="worker6", site="a") worker7 = Worker(name="worker7", site="a") worker8 = Worker(name="worker8", site="a") worker9 = Worker(name="worker9", site="a") worker10 = Worker(name="worker10", site="a") worker11 = Worker(name="worker11", site="a") worker12 = Worker(name="worker12", site="b") worker13 = Worker(name="worker13", site="b") worker14 = Worker(name="worker14", site="b") worker15 = Worker(name="worker15", site="b") worker16 = Worker(name="worker16", site="b") worker17 = Worker(name="worker17", site="b") worker18 = Worker(name="worker18", site="b") worker19 = Worker(name="worker19", site="b") worker20 = Worker(name="worker20", site="b") Rota.add_workers([ worker1, worker2, worker3, worker4, worker5, worker6, worker7, worker8, worker9, worker10, worker11, worker12, worker13, worker14, worker15, worker16, worker17, worker18, worker19, worker20 ]) 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 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), ) for worker in Rota.workers: worker.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), ) for worker in Rota.workers: worker.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_for_all_workers("a", "b", "c") Rota.build_and_solve() Rota.export_rota_to_html("test_triple_shifts_allowed", folder="tests") assert Rota.results.solver.status == "ok" def test_triple_shifts_allowed2(long_rota): Rota = long_rota worker2 = Worker(name="worker2", site="site1", grade=1) worker3 = Worker(name="worker3", site="site1", grade=1) Rota.add_worker(worker2) Rota.add_worker(worker3) Rota.add_shifts( SingleShift(sites=["site1"], name="a", length=8, days=days[2:4], workers_required=1, balance_offset=1), SingleShift(sites=["site1"], name="b", length=8, days=days[2:4], workers_required=1, balance_offset=1), SingleShift(sites=["site1"], name="c", length=8, days=days[2:4], workers_required=1, balance_offset=1), ) worker2.allow_shifts_together("a", "b", "c") worker2.prefer_multi_shift_together = 10 worker3.allow_shifts_together("a", "b", "c") worker3.prefer_multi_shift_together = -10 Rota.build_and_solve() Rota.export_rota_to_html("test_triple_shifts_allowed2", folder="tests") assert Rota.results.solver.status == "ok" def test_multi_shift_prefs(long_rota): Rota = long_rota #Rota.constraint_options["max_shifts_per_week"] = 10 worker2 = Worker(name="worker2", site="site1", grade=1) worker3 = Worker(name="worker3", site="site1", grade=1) worker4 = Worker(name="worker4", site="site1", grade=1) #worker5 = Worker(name="worker5", site="site1", grade=1) Rota.add_worker(worker2) Rota.add_worker(worker3) Rota.add_worker(worker4) #Rota.add_worker(worker5) Rota.add_worker for worker in Rota.workers: worker.allow_shifts_together("a", "b") worker.prefer_multi_shift_together = 0 Rota.add_shifts( SingleShift(sites=["site1"], name="a", length=8, days=days[:1], workers_required=1), SingleShift(sites=["site1"], name="b", length=8, days=days[:1], workers_required=1), SingleShift(sites=["site1"], name="c", length=8, days=days[:1], workers_required=1), SingleShift(sites=["site1"], name="d", length=8, days=days[:1], workers_required=1), SingleShift(sites=["site1"], name="e", length=8, days=days[:1], workers_required=1), #SingleShift(sites=["site1"], name="d", length=8, days=days[:1], workers_required=1), ) Rota.build_and_solve(options={"ratio": 0}) Rota.export_rota_to_html("test_multi_shift_prefs", folder="tests") assert Rota.results.solver.status == "ok" def test_cons_rota(cons_rota): Rota = cons_rota Rota.add_shifts( SingleShift(sites=["a"], name="oncall", length=8, days=days[:5], workers_required=1), SingleShift(sites=["a", "b"], name="evening", length=8, days=days[:5], workers_required=1), ) print("Adding shifts") for worker in Rota.get_workers_by_site("a"): print(worker.name) worker.allow_shifts_together("oncall", "evening") worker.prefer_multi_shift_together = 1 Rota.build_and_solve() Rota.export_rota_to_html("test_cons_rota", folder="tests") assert Rota.results.solver.status == "ok" # Check that everyone at site "a" is assigned 2 or 3 evening shifts for worker in Rota.get_workers_by_site("a"): evening_count = Rota.get_worker_shift_count(worker, "evening") assert 2 <= evening_count <= 3, f"{worker.name} assigned {evening_count} evening shifts" if worker.site == "a": oncall_count = Rota.get_worker_shift_count(worker, "oncall") assert oncall_count == 4 or oncall_count == 5, f"{worker.name} assigned {oncall_count} oncall shifts" # Check that days worked is no more than oncall count days_worked = Rota.get_worker_days_worked(worker) assert days_worked <= oncall_count, f"{worker.name} worked {days_worked} days, more than oncall shifts {oncall_count}" def test_cons_rota2(cons_rota): Rota = cons_rota Rota.add_shifts( SingleShift(sites=["a"], name="oncall", length=8, days=days[:5], workers_required=1), SingleShift(sites=["a", "b"], name="evening", length=8, days=days[:5], workers_required=1), ) print("Adding shifts") for worker in Rota.get_workers_by_site("a"): print(worker.name) worker.allow_shifts_together("oncall", "evening") worker.prefer_multi_shift_together = -1 Rota.build_and_solve() Rota.export_rota_to_html("test_cons_rota", folder="tests") assert Rota.results.solver.status == "ok" # Check that everyone at site "a" is assigned 2 or 3 evening shifts for worker in Rota.get_workers_by_site("a"): evening_count = Rota.get_worker_shift_count(worker, "evening") assert 2 <= evening_count <= 3, f"{worker.name} assigned {evening_count} evening shifts" if worker.site == "a": oncall_count = Rota.get_worker_shift_count(worker, "oncall") assert oncall_count == 4 or oncall_count == 5, f"{worker.name} assigned {oncall_count} oncall shifts" # Check that days worked is now more than oncall count days_worked = Rota.get_worker_days_worked(worker) assert days_worked > oncall_count, f"{worker.name} worked {days_worked} days, more than oncall shifts {oncall_count}" def test_cons_rota3(cons_rota): Rota = cons_rota Rota.terminate_on_warning.remove("Worker/no valid shifts") Rota.add_shifts( #SingleShift(sites=["a"], name="oncall weekday", length=8, days=days[:4], workers_required=1, # constraint=[ # {"name": "pre", "options": 2}, # {"name": "post", "options": 2}, # ], # ), ##SingleShift(sites=["a"], name="oncall weekend", length=8, days=days[4:], workers_required=1), #SingleShift(sites=["a", "b"], name="evening", length=8, days=days[:4], workers_required=1, # constraint=[ # {"name": "pre", "options": 2}, # {"name": "post", "options": 2}, # ], # ), SingleShift(sites=["a"], name="help", length=8, days=days[4:], workers_required=1, #assign_as_block=True, #force_as_block=True #constraint=[ # {"name": "pre", "options": 2}, # {"name": "post", "options": 2}, #], ) , SingleShift(sites=["b"], name="weekend b", length=8, days=days[5:], workers_required=1, #constraint=[ # {"name": "pre", "options": 2}, # {"name": "post", "options": 2}, #], ), ) #print("Adding shifts") #for worker in Rota.get_workers_by_site("a"): # print(worker.name) ## worker.allow_shifts_together("oncall weekday", "evening") ## #worker.allow_shifts_together("oncall weekend", "weekend") ## worker.prefer_multi_shift_together = 1 # worker.assign_as_block_preferences = {"help":10} #for worker in Rota.get_workers_by_site("b"): # worker.assign_as_block_preferences = {"weekend b":-10} Rota.build_and_solve() Rota.export_rota_to_html("test_cons_rota", folder="tests") assert Rota.results.solver.status == "ok"