diff --git a/gen_cons.py b/gen_cons.py index dcb4e21..0a32e7f 100644 --- a/gen_cons.py +++ b/gen_cons.py @@ -121,6 +121,14 @@ def main( constraint=[], ), ) + Rota.allow_shifts_together( + + "oncall weekday", + "Oncall weekend", + "weekend a", + "evening", + + ) # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"]) diff --git a/output/timetable.css b/output/timetable.css index e9d27c8..0adbba2 100644 --- a/output/timetable.css +++ b/output/timetable.css @@ -321,4 +321,8 @@ table.transposed th.bank-holiday { .locum-shift { background-color: lightgreen; +} + +.multi-shift { + line-height: 1; } \ No newline at end of file diff --git a/output/timetable.js b/output/timetable.js index b240ee1..2d23acd 100644 --- a/output/timetable.js +++ b/output/timetable.js @@ -295,7 +295,10 @@ $(".table-div .worker-row .worker").each((n, tr) => { locum_shift_counts = jtr.data("locum-shift-counts") shift_targets = jtr.data("worker-targets") + console.log(shift_counts, locum_shift_counts, shift_targets) + oshifts.forEach((s) => { + console.log(s, shift_counts, shift_targets, locum_shift_counts) if (s in shift_targets && shift_targets[s] > 0) { if (s in shift_counts) { c = shift_counts[s]; @@ -426,6 +429,7 @@ $(".table-div + pre").each((n, pre) => { }) + $("td").hover((e) => { // start hover worker_td = $(e.target).closest('tr').find('td:first')[0] @@ -436,31 +440,37 @@ $("td").hover((e) => { $($(el).closest('tr').find('td').get(index)).addClass("pair-match") }) } - if (e.target.dataset.shift != undefined && e.target.dataset.shift.length > 0) { - // Highlight all cells containing this shift, even if multiple shifts per day - $("td").filter(function() { - let shifts = $(this).attr("data-shift"); - if (!shifts) return false; - return shifts.split(",").map(s => s.trim()).includes(e.target.dataset.shift); - }).addClass("shift-highlight"); - // For remote site, similar logic if needed + // Find the span under the mouse, if any + let targetSpan = $(e.target).closest("span.multi-shift-shift").get(0) || + ($(e.target).hasClass("multi-shift-shift") ? e.target : null); + + if (targetSpan) { + // Only highlight spans with the same text + let shiftDisplay = $(targetSpan).text().trim(); + $("span.multi-shift-shift").filter(function() { + return $(this).text().trim() === shiftDisplay; + }).addClass("shift-highlight"); + } else if (e.target.dataset.shift != undefined && e.target.dataset.shift.length > 0) { + // Fallback: highlight all cells containing ANY of the hovered shifts (legacy) + let hovered_shifts = e.target.dataset.shift.split(",").map(s => s.trim()); $("td").filter(function() { let shifts = $(this).attr("data-shift"); - let remote = $(this).attr("data-remote-site"); if (!shifts) return false; - return shifts.split(",").map(s => s.trim()).includes(e.target.dataset.shift) && - remote == e.target.dataset.shiftRemoteSite; - }).addClass("remote-site-match"); + let cell_shifts = shifts.split(",").map(s => s.trim()); + return hovered_shifts.some(s => cell_shifts.includes(s)); + }).addClass("shift-highlight"); } }, (e) => { // end hover $("td.pair-match").removeClass("pair-match") $("td.shift-highlight").removeClass("shift-highlight") + $("span.shift-highlight").removeClass("shift-highlight") $("td.remote-site-match").removeClass("remote-site-match") }); + $("#rota-table tr td:first-child").each((n, td) => { $(td).click(() => { viewWorker(td); diff --git a/rota/shifts.py b/rota/shifts.py index fc764ef..b50ff7a 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -257,7 +257,7 @@ class RotaBuilder(object): self.SHIFT_BOUNDS = SHIFT_BOUNDS self.shifts: List[SingleShift] = [] # type List[SingleShift] - self.allowed_multi_shift_sets: list[set[str]] = [] + #self.allowed_multi_shift_sets: list[set[str]] = [] self.use_previous_shifts = use_previous_shifts self.use_shift_balance_extra = use_shift_balance_extra @@ -603,6 +603,16 @@ class RotaBuilder(object): initialize=0, ) + self.model.multi_shift_together = Var( + ((worker.id, week, day, idx) + for worker in self.workers + for week, day in self.get_week_day_combinations() + for idx, allowed_set in enumerate(getattr(worker, "allowed_multi_shift_sets", [])) + if allowed_set.issubset(self.get_shift_names_by_week_day(week, day))), + domain=Binary, + initialize=0, + ) + # Used to limit number of workers on night shift per site # if we force a binary it will in effect hard constrain to < 2 # from a single site @@ -1398,6 +1408,25 @@ class RotaBuilder(object): ) logging.debug(f"Generate worker constraints: {worker.name}") + for week, day in self.get_week_day_combinations(): + shifts_today = self.get_shift_names_by_week_day(week, day) + worker_allowed_sets = getattr(worker, "allowed_multi_shift_sets", []) + #for idx, allowed_set in enumerate(worker_allowed_sets): + # if allowed_set.issubset(shifts_today): + # var = self.model.multi_shift_together[worker.id, week, day, idx] + # self.model.constraints.add( + # sum(self.model.works[worker.id, week, day, shift] for shift in allowed_set) + # == len(allowed_set) * var + # ) + for idx, allowed_set in enumerate(worker_allowed_sets): + if allowed_set.issubset(shifts_today): + var = self.model.multi_shift_together[worker.id, week, day, idx] + # var is 1 iff all shifts in allowed_set are assigned + for shift in allowed_set: + self.model.constraints.add(var <= self.model.works[worker.id, week, day, shift]) + self.model.constraints.add( + var >= sum(self.model.works[worker.id, week, day, shift] for shift in allowed_set) - len(allowed_set) + 1 + ) if self.get_workers_who_require_locums(): @@ -2416,55 +2445,40 @@ class RotaBuilder(object): # single shift per day (unless multi-shift allowed) # This is signifantly slower so only enable if required - if self.allowed_multi_shift_sets: - shifts_today = self.get_shift_names_by_week_day(week, day) - if shifts_today: - # Build all allowed sets that are subsets of today's shifts - allowed_sets_today = [ - allowed_set for allowed_set in self.allowed_multi_shift_sets - if allowed_set.issubset(shifts_today) - ] + shifts_today = self.get_shift_names_by_week_day(week, day) + assigned_vars = [self.model.works[worker.id, week, day, shift] for shift in shifts_today] + worker_allowed_sets = getattr(worker, "allowed_multi_shift_sets", []) + allowed_sets_today = [ + allowed_set for allowed_set in worker_allowed_sets + if allowed_set.issubset(shifts_today) + ] - # For each possible assignment, only allow: - # - a single shift - # - OR exactly one of the allowed sets - assigned_vars = [self.model.works[worker.id, week, day, shift] for shift in shifts_today] - # The sum of assigned shifts must be 1 (single shift) or exactly len(allowed_set) for one allowed set - allowed_sums = [1] + [len(allowed_set) for allowed_set in allowed_sets_today] - # Enforce: sum in allowed_sums + if allowed_sets_today: + # Forbid any multi-shift assignment that includes shifts from different allowed sets + # Build all possible multi-shift assignments + all_multi = [set(c) for i in range(2, len(shifts_today)+1) for c in itertools.combinations(shifts_today, i)] + # Build all allowed multi-shift patterns (all non-empty subsets of each allowed set) + allowed_patterns = [{shift} for shift in shifts_today] + for allowed_set in allowed_sets_today: + allowed_patterns += [set(s) for i in range(2, len(allowed_set)+1) for s in itertools.combinations(allowed_set, i)] + # Forbid all other multi-shift combinations + forbidden = [s for s in all_multi if s not in allowed_patterns] + for forbidden_set in forbidden: self.model.constraints.add( - sum(assigned_vars) <= max(allowed_sums) + sum(self.model.works[worker.id, week, day, shift] for shift in forbidden_set) + <= len(forbidden_set) - 1 ) - self.model.constraints.add( - sum(assigned_vars) >= min(allowed_sums) - ) - # Now, for each allowed set, prevent partial assignment of the set - for allowed_set in allowed_sets_today: - # If any shift in allowed_set is assigned, all must be assigned - for shift in allowed_set: - self.model.constraints.add( - sum(self.model.works[worker.id, week, day, s] for s in allowed_set) - >= len(allowed_set) * self.model.works[worker.id, week, day, shift] - ) - # No other shifts can be assigned together with the allowed set - for shift in shifts_today: - if shift not in allowed_set: - self.model.constraints.add( - sum(self.model.works[worker.id, week, day, s] for s in allowed_set) - + self.model.works[worker.id, week, day, shift] - <= len(allowed_set) - ) - - - # Otherwise only one shift per day - else: + # Still, at most all shifts per day self.model.constraints.add( - 1 - >= sum( - self.model.works[worker.id, week, day, shift] - for shift in self.get_shift_names_by_week_day(week, day) - ) + sum(self.model.works[worker.id, week, day, shift] for shift in shifts_today) + <= len(shifts_today) ) + else: + # Only one shift per day if no allowed sets + self.model.constraints.add( + sum(assigned_vars) <= 1 + ) + # This applies to locums as well if self.get_workers_who_require_locums(): self.model.constraints.add( @@ -2647,6 +2661,14 @@ class RotaBuilder(object): # Define an objective function with model as input, to pass later def obj_rule(m): # c = len(workers) + prefer_multi_shift_expr = sum( + getattr(worker, "prefer_multi_shift_together", 0) * + self.model.multi_shift_together[worker.id, week, day, idx] + for worker in self.workers + for week, day in self.get_week_day_combinations() + for idx, allowed_set in enumerate(getattr(worker, "allowed_multi_shift_sets", [])) + if allowed_set.issubset(self.get_shift_names_by_week_day(week, day)) + ) balance_modifier_constant = 10 balance_quadratic_shift_modifier_constant = 3 @@ -2840,16 +2862,17 @@ class RotaBuilder(object): - work_requests + locum_shift_balancing + true_quadratic_shift_balancing + - prefer_multi_shift_expr ) # add objective function to the model. rule (pass function) or expr (pass expression directly) self.model.obj = Objective(rule=obj_rule, sense=minimize) - def allow_shifts_together(self, *shift_names: str): - """ - Allow the given set of shift names to be assigned together on a single day. - """ - self.allowed_multi_shift_sets.append(set(shift_names)) + #def allow_shifts_together(self, *shift_names: str): + # """ + # Allow the given set of shift names to be assigned together on a single day. + # """ + # self.allowed_multi_shift_sets.append(set(shift_names)) def add_warning(self, warning_type: str, message: str): print(f"[bold red]WARNING:[/bold red] {warning_type} - {message}") @@ -3032,6 +3055,13 @@ class RotaBuilder(object): for n in range(int((end_date - start_date).days) + 1): yield start_date + datetime.timedelta(n) + def allow_shifts_together_for_all_workers(self, *shift_names: str): + """ + Allow the given set of shift names to be assigned together on a single day for all workers. + """ + for worker in self.workers: + worker.allowed_multi_shift_sets.append(frozenset(shift_names)) + def build_shifts(self): """ Process the added shifts @@ -3151,6 +3181,16 @@ class RotaBuilder(object): for shift in self.get_shifts_with_constraint("post"): self.max_post = max(shift.constraint_options["post"], self.max_post) + # --- Check allowed_multi_shift_sets validity --- + all_shift_names = set(s.name for s in self.shifts) + for worker in self.workers: + for allowed_set in getattr(worker, "allowed_multi_shift_sets", []): + missing = set(allowed_set) - all_shift_names + if missing: + raise InvalidShift( + f"Worker '{worker.name}' has allowed_multi_shift_set(s) with non-existent shift(s): {missing}" + ) + ## For paired shifts we create proxy shifts # self.paired_shift_map = {} # for shifts in self.paired_shifts: @@ -3414,6 +3454,18 @@ class RotaBuilder(object): return group_workers + def get_workers_by_site(self, site: str) -> list[Worker]: + """ + Returns a list of workers assigned to a specific site. + + Args: + site (str): The site name. + + Returns: + list[Worker]: List of workers at the given site. + """ + return [worker for worker in self.workers if worker.site == site] + def get_workers_by_grade(self) -> Dict[int, list[Worker]]: group_workers: dict[int, list[Worker]] = defaultdict(list) for worker in self.workers: @@ -3466,7 +3518,7 @@ class RotaBuilder(object): # RESULTS def export_rota_to_html( - self, filename: str = "rota", folder=None, timestamp_filename=True + self, filename: str = "rota", folder=None, timestamp_filename=False ): if timestamp_filename: filename = f"{filename}_{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}" @@ -3562,15 +3614,15 @@ class RotaBuilder(object): shifts = [] w = [f"{worker.name:20}"] for week, day in self.get_week_day_combinations(): - a = "-" + shift_display_char = "-" shift_names = [] for shift in self.get_shift_names_by_week_day(week, day): if model.works[worker.id, week, day, shift].value > 0.5: shift_names.append(self.get_shift_by_name(shift).get_display_char()) if shift_names: - a = "".join(shift_names) + shift_display_char = "".join(shift_names) shifts.extend(shift_names) - w.append(a) + w.append(shift_display_char) shift_count = "" for s in set(shifts): shift_count = shift_count + f"{s}: {shifts.count(s)}, " @@ -3676,7 +3728,7 @@ class RotaBuilder(object): for week, day in self.get_week_day_combinations(): d = self.start_date + datetime.timedelta(n) n = n + 1 - a = "-" + shift_display_char = "-" shift_names = [] assigned_shift_names = [] locum = False @@ -3693,7 +3745,7 @@ class RotaBuilder(object): assigned_shift_names.append(shift) locum = True if shift_names: - a = "
".join(shift_names) + shift_display_char = "
".join([""+name+"" for name in shift_names]) shifts.extend(assigned_shift_names) shift_name = assigned_shift_names[0] if assigned_shift_names else "" @@ -3704,7 +3756,6 @@ class RotaBuilder(object): css_class += " multi-shift" title = f"{shift_name} ({d})" - css_class = day unavailable_reason = "" if model.available[worker.id, week, day] > 0: available = True @@ -3751,7 +3802,13 @@ class RotaBuilder(object): css_class = " ".join((css_class, "night-shift")) shift_tds.append( - f"{a}" + f"{shift_display_char}" ) shift_count = "" @@ -3961,7 +4018,7 @@ class RotaBuilder(object): temp = "" for shift in self.get_shift_names_by_week_day(week, day): - if self.model.works[worker.id, week, day, shift].value > 0: + if self.model.works[worker.id, week, day, shift].value > 0.5: shifts[d] = shift return shifts @@ -3994,6 +4051,27 @@ class RotaBuilder(object): return shifts + def get_worker_shift_count(self, worker: Worker, shift_name:str) -> int: + """Returns the number of shifts of a specific type assigned to a worker.""" + count = 0 + for week, day in self.get_week_day_combinations(): + if shift_name in self.get_shift_names_by_week_day(week, day): + if self.model.works[worker.id, week, day, shift_name].value > 0.5: + count += 1 + return count + + def get_worker_days_worked(self, worker: Worker) -> int: + """ + Returns the number of unique days on which the worker is assigned at least one shift. + """ + days_worked = 0 + for week, day in self.get_week_day_combinations(): + for shift in self.get_shift_names_by_week_day(week, day): + if self.model.works[worker.id, week, day, shift].value > 0.5: + days_worked += 1 + break # Only count each day once, even if multiple shifts + return days_worked + def get_worker_shift_list_string(self, worker: Worker) -> str: shifts = self.get_worker_shift_list(worker) diff --git a/rota/workers.py b/rota/workers.py index fac07ee..682e6db 100644 --- a/rota/workers.py +++ b/rota/workers.py @@ -111,6 +111,9 @@ class Worker(BaseModel): bank_holiday_extra: int = 0 pair: int | str | None = None locum: bool = False + allowed_multi_shift_sets: list[set] = [] # or list/set of frozensets + prefer_multi_shift_together: int = 0 # Default: no preference + # Set to a positive integer to prefer, negative to discourage, 0 for neutral shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist @@ -389,3 +392,8 @@ class Worker(BaseModel): if shift in self.fte_adj_shifts: return self.fte_adj_shifts[shift] return self.fte_adj + + def allow_shifts_together(self, *shifts): + if not hasattr(self, "allowed_multi_shift_sets"): + self.allowed_multi_shift_sets = [] + self.allowed_multi_shift_sets.append(frozenset(shifts)) \ No newline at end of file diff --git a/test/test_multiple_shifts_per_day.py b/test/test_multiple_shifts_per_day.py index 5f35e4d..b4a36f6 100644 --- a/test/test_multiple_shifts_per_day.py +++ b/test/test_multiple_shifts_per_day.py @@ -12,6 +12,50 @@ def basic_rota(): 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( @@ -21,12 +65,6 @@ def test_single_shift_per_day_default(basic_rota): Rota.build_and_solve() assert Rota.results.solver.status == "error" # Only one shift per day should be assigned - 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) <= 1 def test_allow_two_shifts_per_day(basic_rota): Rota = basic_rota @@ -34,7 +72,8 @@ def test_allow_two_shifts_per_day(basic_rota): 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), ) - Rota.allow_shifts_together("a", "b") + 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" @@ -58,7 +97,8 @@ def test_only_specified_shifts_can_be_double_assigned(basic_rota): 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), ) - Rota.allow_shifts_together("a", "b") + 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" @@ -79,8 +119,111 @@ def test_triple_shifts_allowed(basic_rota): 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("a", "b", "c") + 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}"