From 8402cc92fda81b18286c96138a871849ea53c809 Mon Sep 17 00:00:00 2001 From: Ross Date: Sun, 22 Jun 2025 10:53:17 +0100 Subject: [PATCH] fix bug with pre / post shift constraints --- rota/shifts.py | 45 +---- test/test_rota_clear.py | 371 +++++++++++++++++++++++++++++++++++++++- test/test_workers.py | 5 +- 3 files changed, 381 insertions(+), 40 deletions(-) diff --git a/rota/shifts.py b/rota/shifts.py index 808c478..3ada902 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -2626,10 +2626,9 @@ class RotaBuilder(object): ): for n in range(0, constraint_shift.constraint_options["pre"]): if day in constraint_shift.days: - # week_start_date = self.get_week_start_date(week) - # try: - # if constraint_shift.start_date <= week_start_date < constraint_shift.end_date: - # works = 1 + # Only apply if worker can work this shift + if worker.site not in constraint_shift.sites: + continue try: works = self.model.works[ worker.id, week, day, constraint_shift.name @@ -2649,50 +2648,23 @@ class RotaBuilder(object): ) if shiftname != constraint_shift.name for w in workers + if w.site in constraint_shift.sites # Only workers who can work this shift ) ) - # for constraint_shift in self.get_shifts_with_constraint( "pre"): - # to_sum = [] - # for n in range( - # 0, constraint_shift.constraint_options["pre"] - # ): - # for w in workers: - # for wk, dy, shiftname in self.get_week_day_shift_names_by_week_day( - # pre_map[n][1], pre_map[n][2] - # ): - # if shiftname != constraint_shift.name: - # to_sum.append(pre_map[n][0] - # * self.model.works[ - # w.id, wk, dy, shiftname - # ]) - - # if day in constraint_shift.days: - # self.model.constraints.add( - # 1 - # >= self.model.works[ - # worker.id, week, day, constraint_shift.name - # ] - # + sum( - # to_sum - # ) - # ) for constraint_shift in self.get_shifts_with_constraints( "post" # , week=week ): for n in range(0, constraint_shift.constraint_options["post"]): if day in constraint_shift.days: - # week_start_date = self.get_week_start_date(week) - # if constraint_shift.start_date <= week_start_date < constraint_shift.end_date: - # works = 1 - # else: + # Only apply if worker can work this shift + if worker.site not in constraint_shift.sites: + continue try: works = self.model.works[ worker.id, week, day, constraint_shift.name ] except KeyError: - # If we can't find the key the shift is not assigned - # so we don't need to worry about it continue self.model.constraints.add( 1 @@ -2707,10 +2679,9 @@ class RotaBuilder(object): ) if shiftname != constraint_shift.name for w in workers + if w.site in constraint_shift.sites # Only workers who can work this shift ) ) - # except KeyError: - # pass # Night constraint means we won't assign a shift the day before # an unavailability diff --git a/test/test_rota_clear.py b/test/test_rota_clear.py index 1f616cc..edc0d2a 100644 --- a/test/test_rota_clear.py +++ b/test/test_rota_clear.py @@ -3,6 +3,27 @@ from rota.shifts import RotaBuilder, SingleShift, days from rota.workers import Worker import datetime + +def generate_basic_rota( + weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7) +): + Rota = RotaBuilder( + start_date, + weeks_to_rota=weeks_to_rota, + ) + + # Add a few workers + if workers: + Rota.add_workers( + [ + Worker(name=f"worker{i:02}", site="group1", grade=1) + for i in range(1, workers + 1) + ] + ) + + return Rota + + @pytest.fixture def demo_rota_clear(): weeks_to_rota = 10 @@ -28,6 +49,7 @@ def demo_rota_clear(): Rota.add_workers(workers) return Rota + def test_pre(demo_rota_clear): Rota = demo_rota_clear Rota.add_shifts( @@ -82,4 +104,351 @@ def test_pre(demo_rota_clear): # assert "a-c" not in shifts_string # assert "c-a" not in shifts_string # assert "b-c" not in shifts_string - # assert "c-b" not in shifts_string \ No newline at end of file + # assert "c-b" not in shifts_string + + +def test_pre_post_clear(demo_rota_clear): + Rota = generate_basic_rota(workers=4, weeks_to_rota=16) + + workers = Rota.get_workers() + + # for worker in workers: + # worker.add_force_assign_with("a", "b") + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[:5], + workers_required=2, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 1, + }, + {"name": "post", "options": 1}, + ], + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[5:], + workers_required=2, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 1, + }, + {"name": "post", "options": 1}, + ], + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.export_rota_to_html("test_pre_post_clear", folder="tests") + + assert Rota.results.solver.status == "ok" + for worker in Rota.workers: + shifts = Rota.get_worker_shift_list(worker) + shifts_string = "".join([i if i != "" else "-" for i in shifts]) + + # check for patterns that should not occur + assert "ab" not in shifts_string + assert "ba" not in shifts_string + # assert "a-b" not in shifts_string + # assert "b-a" not in shifts_string + + +def test_pre_post_clear2(demo_rota_clear): + Rota = generate_basic_rota(workers=4, weeks_to_rota=16) + + workers = Rota.get_workers() + + # for worker in workers: + # worker.add_force_assign_with("a", "b") + + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[:5], + workers_required=2, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 2, + }, + {"name": "post", "options": 2}, + ], + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[5:], + workers_required=2, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 2, + }, + {"name": "post", "options": 2}, + ], + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.export_rota_to_html("test_pre_post_clear", folder="tests") + + assert Rota.results.solver.status == "ok" + for worker in Rota.workers: + shifts = Rota.get_worker_shift_list(worker) + shifts_string = "".join([i if i != "" else "-" for i in shifts]) + + # check for patterns that should not occur + assert "ab" not in shifts_string + assert "ba" not in shifts_string + assert "a-b" not in shifts_string + assert "b-a" not in shifts_string + + +def test_pre_post_clear_multi_group(demo_rota_clear): + Rota = generate_basic_rota(workers=2, weeks_to_rota=16) + + workers = Rota.get_workers() + + Rota.add_workers( + ( + Worker(name="worker3", site="group2", grade=1), + Worker(name="worker4", site="group2", grade=1), + ) + ) + + Rota.add_shifts( + SingleShift( + sites=("group1", "group2"), + name="a", + length=8, + days=days[:5], + workers_required=2, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 2, + }, + {"name": "post", "options": 2}, + ], + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[5:], + workers_required=1, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 2, + }, + {"name": "post", "options": 2}, + ], + ), + SingleShift( + sites=("group2",), + name="c", + length=8, + days=days[5:], + workers_required=1, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 2, + }, + {"name": "post", "options": 2}, + ], + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.export_rota_to_html("test_pre_post_clear", folder="tests") + + assert Rota.results.solver.status == "ok" + for worker in Rota.workers: + shifts = Rota.get_worker_shift_list(worker) + shifts_string = "".join([i if i != "" else "-" for i in shifts]) + + # check for patterns that should not occur + assert "ab" not in shifts_string + assert "ba" not in shifts_string + assert "a-b" not in shifts_string + assert "b-a" not in shifts_string + + assert "ac" not in shifts_string + assert "ca" not in shifts_string + assert "a-c" not in shifts_string + assert "c-a" not in shifts_string + +def test_pre_post_clear_force_assign(demo_rota_clear): + Rota = generate_basic_rota(workers=6, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_force_assign_with("a", "o") + + Rota.add_workers( + ( + Worker(name="workera", site="group2", grade=1), + Worker(name="workerb", site="group2", grade=1), + ) + ) + + Rota.add_shifts( + SingleShift( + sites=("group1", "group2"), + name="a", + length=8, + days=days[:5], + workers_required=1, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 1, + }, + {"name": "post", "options": 1}, + ], + ), + SingleShift( + sites=("group1",), + name="o", + length=8, + days=days[:7], + workers_required=1, + assign_as_block=False, # global setting is off + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[5:], + workers_required=1, + assign_as_block=False, # global setting is off + #constraint=[ + # { + # "name": "pre", + # "options": 2, + # }, + # {"name": "post", "options": 2}, + #], + ), + SingleShift( + sites=("group2",), + name="c", + length=8, + days=days[5:], + workers_required=1, + assign_as_block=False, # global setting is off + #constraint=[ + # { + # "name": "pre", + # "options": 1, + # }, + # {"name": "post", "options": 1}, + #], + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.export_rota_to_html("test_pre_post_clear", folder="tests") + + assert Rota.results.solver.status == "ok" + +def test_pre_post_clear_force_assign2(demo_rota_clear): + Rota = generate_basic_rota(workers=6, weeks_to_rota=16) + + workers = Rota.get_workers() + + for worker in workers: + worker.add_force_assign_with("a", "o") + worker.add_force_assign_with("b", "o") + + Rota.add_workers( + ( + Worker(name="workera", site="group2", grade=1), + Worker(name="workerb", site="group2", grade=1), + ) + ) + + Rota.add_shifts( + SingleShift( + sites=("group1", "group2"), + name="a", + length=8, + days=days[:5], + workers_required=1, + assign_as_block=False, # global setting is off + constraint=[ + #{ + # "name": "pre", + # "options": 1, + #}, + #{"name": "post", "options": 1}, + ], + ), + SingleShift( + sites=("group1",), + name="o", + length=8, + days=days[:7], + workers_required=1, + assign_as_block=False, # global setting is off + ), + SingleShift( + sites=("group1",), + name="b", + length=8, + days=days[5:], + workers_required=1, + assign_as_block=False, # global setting is off + constraint=[ + #{ + # "name": "pre", + # "options": 2, + #}, + #{"name": "post", "options": 2}, + ], + ), + SingleShift( + sites=("group2",), + name="c", + length=8, + days=days[5:], + workers_required=1, + assign_as_block=False, # global setting is off + constraint=[ + { + "name": "pre", + "options": 1, + }, + {"name": "post", "options": 1}, + ], + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + Rota.export_rota_to_html("test_pre_post_clear", folder="tests") + + assert Rota.results.solver.status == "ok" + for worker in Rota.workers: + shifts = Rota.get_worker_shift_list(worker) + shifts_string = "".join([i if i != "" else "-" for i in shifts]) + + # check for patterns that should not occur + assert "ac" not in shifts_string + assert "ca" not in shifts_string \ No newline at end of file diff --git a/test/test_workers.py b/test/test_workers.py index 33b4895..87055e9 100644 --- a/test/test_workers.py +++ b/test/test_workers.py @@ -964,10 +964,11 @@ def test_worker_assign_as_block_preference_multiple_shifts(): ), ) - shift_list_1 = Rota.get_worker_shift_list_string(worker1) - assert shift_list_1.count("aaaa") == 3, "Worker 1 should have all 'a' shifts grouped together (block)" Rota.build_and_solve(options={"ratio": 0.0}) + + shift_list_1 = Rota.get_worker_shift_list_string(worker1) + assert shift_list_1.count("aaaa") == 3, "Worker 1 should have all 'a' shifts grouped together (block)" Rota.export_rota_to_html("test_worker_assign_as_block_preference", folder="tests") def test_worker_assign_as_block_preference_multiple_sites():