latest rotas

gen.py

@@ -19,8 +19,6 @@ sites = (
     "proc",
 )
 
-NIGHT_REGISTRAR_NUMBER = 4
-
 from rota.workers import (
     Worker,
     NotAvailableToWork,
@@ -35,11 +33,11 @@ from rota.workers import (
 def main(
     suspend: bool = False,
     solve: bool = True,
-    time_to_run: int = 60 * 60,
+    time_to_run: int = 60 * 60 * 4,
     ratio: float = 0.001,
-    start_date: datetime.datetime = "2025-03-03",
-    weeks: int = 26,
-    bom: int = 1,
+    start_date: datetime.datetime = "2025-09-01",
+    weeks: int = 16,
+    bom: int = 2,
 ):
     rota_start_date = start_date.date()
     suspend_on_finish = suspend
@@ -60,16 +58,16 @@ def main(
     Rota.constraint_options["hard_constrain_pair_separation"] = True
     # Rota.constraint_options["avoid_st2_first_month"] = True
 
-    wr = [
-        WorkerRequirement(
-            end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
-            number=3,
-        ),
-        WorkerRequirement(
-            start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
-            number=4,
-        ),
-    ]
+    #wr = [
+    #    WorkerRequirement(
+    #        end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+    #        number=3,
+    #    ),
+    #    WorkerRequirement(
+    #        start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+    #        number=4,
+    #    ),
+    #]
 
     Rota.add_shifts(
         SingleShift(
@@ -226,7 +224,7 @@ def main(
             balance_offset=3.9,
             balance_weighting=1,
             # hard_constrain_shift=False,
-            workers_required=wr,
+            workers_required=4,
             force_as_block=True,
             rota_on_nwds=True,
             constraint=[
@@ -256,7 +254,7 @@ def main(
             balance_offset=2.9,
             balance_weighting=1,
             # hard_constrain_shift=False,
-            workers_required=wr,
+            workers_required=4,
             force_as_block=True,
             rota_on_nwds=True,
             constraint=[
@@ -267,68 +265,20 @@ def main(
                     "name": "require_remote_site_presence_week",
                     "options": ("plymouth", 1),
                 },
-                {"name": "limit_grade_number", "options": {2: 1}},
+                {"name": "limit_grade_number", "options": {2: 2}},
                 {"name": "minimum_grade_number", "options": (4, 1)},
             ],
             #end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
         ),
-        # SingleShift(
-        #    sites=sites,
-        #    name="night_weekday4",
-        #    length=12.25,
-        #    days=days[:4],
-        #    balance_offset=3.9,
-        #    balance_weighting=1,
-        #    # hard_constrain_shift=False,
-        #    workers_required=NIGHT_REGISTRAR_NUMBER,
-        #    force_as_block=True,
-        #    rota_on_nwds=True,
-        #    constraint=[
-        #        {"name": "night"},
-        #        {"name": "pre", "options": 2},
-        #        {"name": "post", "options": 2},
-        #        {
-        #            "name": "require_remote_site_presence_week",
-        #            "options": ("plymouth", 1),
-        #        },
-        #        {"name": "limit_grade_number", "options": {2: 1}},
-        #        {"name": "minimum_grade_number", "options": (4, 1)},
-        #    ],
-        #    start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
-        # ),
-        # SingleShift(
-        #    sites=[
-        #        *sites,
-        #        "plymouth twilights, weekends and weekend nights",
-        #        "truro twilights and weekend nights",
-        #        "weekend nights",
-        #        "truro twilights, weekends and weekend nights",
-        #    ],
-        #    name="night_weekend4",
-        #    length=12.25,
-        #    days=days[4:],
-        #    balance_offset=2.9,
-        #    balance_weighting=1,
-        #    # hard_constrain_shift=False,
-        #    workers_required=NIGHT_REGISTRAR_NUMBER,
-        #    force_as_block=True,
-        #    rota_on_nwds=True,
-        #    constraint=[
-        #        {"name": "night"},
-        #        {"name": "pre", "options": 2},
-        #        {"name": "post", "options": 3},
-        #        {
-        #            "name": "require_remote_site_presence_week",
-        #            "options": ("plymouth", 1),
-        #        },
-        #        {"name": "limit_grade_number", "options": {2: 1}},
-        #        {"name": "minimum_grade_number", "options": (4, 1)},
-        #    ],
-        #    start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
-        # ),
     )
 
-    # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"])
+    # Prevent shifts for ST2s for 2 weeks
+    Rota.add_grade_constraint_by_week([2], [1, 2], [shift.name for shift in Rota.shifts])
+
+    Rota.add_min_summed_grade_by_shifts_per_day_constraint(["weekend_plymouth1", "weekend_plymouth2"], 6)
+    Rota.add_min_summed_grade_by_shifts_per_day_constraint(["plymouth_twilight", "plymouth_bank_holidays"], 6)
+
+    Rota.terminate_on_warning.remove("Worker/no valid shifts")
 
     load_leave = True
     Rota.build_shifts()
@@ -341,10 +291,17 @@ def main(
 
         n = 0
         for worker in workers:
-            n = n + 1
             worker_name = worker
+
+            if worker_name == "Name":
+                continue
+            n = n + 1
             w = workers[worker]
-            site = w["site"]
+            try:
+                site = w["site"]
+            except KeyError:
+                print(f"Worker {worker} has no site")
+                raise KeyError
             grade = w["grade"]
             try:
                 fte = float(w["fte"]) * 100
@@ -466,22 +423,26 @@ def main(
 
             shift_fte_overrides = {}
 
-            if worker_name == "Ben Kemp":
+            #if worker_name == "Ben Kemp":
+            #    shift_fte_overrides = {
+            #        "plymouth_twilight": 100,
+            #        "weekend_plymouth1": 50,
+            #        "weekend_plymouth2": 50,
+            #    }
+            #elif worker_name == "Joel Lim":
+            #    shift_fte_overrides = {
+            #        "plymouth_twilight": 100,
+            #        "weekend_exeter": 50,
+            #    }
+            if worker_name == "Nang Thiriphoo":
                 shift_fte_overrides = {
-                    "plymouth_twilight": 100,
-                    "weekend_plymouth1": 50,
-                    "weekend_plymouth2": 50,
-                }
-            elif worker_name == "Joel Lim":
-                shift_fte_overrides = {
-                    "plymouth_twilight": 100,
-                    "weekend_exeter": 50,
+                    "weekend_exeter": 40,
                 }
 
             w = Worker(
                 name=worker_name,
                 site=site.lower(),
-                grade=int(grade[2]),
+                grade=int(grade),
                 id=n,
                 fte=int(fte),
                 nwds=nwds,

gen_cons.py

@@ -0,0 +1,234 @@
+import datetime
+import os
+import sys
+import time
+from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WorkerRequirement, days
+import typer
+import subprocess
+
+from rota.workers import (
+    Worker,
+    NotAvailableToWork,
+    NonWorkingDays,
+    WorkRequests,
+    PreferenceNotToWork,
+    OutOfProgramme,
+)
+
+
+app = typer.Typer()
+
+
+sites = ("rota a", "rota b", "rota c")
+
+
+@app.command()
+def main(
+    suspend: bool = False,
+    solve: bool = True,
+    time_to_run: int = 60 * 60,
+    ratio: float = 0.001,
+    start_date: datetime.datetime = "2025-09-01",
+    weeks: int = 22,
+    bom: int = 1,
+):
+    rota_start_date = start_date.date()
+    suspend_on_finish = suspend
+
+    Rota = RotaBuilder(
+        rota_start_date,
+        weeks_to_rota=weeks,
+        balance_offset_modifier=bom,
+        name="cons rota",
+    )
+    # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
+    Rota.constraint_options["balance_weekends"] = True
+    Rota.constraint_options["max_weekend_frequency"] = 4
+    Rota.constraint_options["max_shifts_per_week"] = 20
+    # Rota.constraint_options["avoid_st2_first_month"] = True
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("rota a",),
+            name="oncall",
+            length=12.5,
+            days=days,
+            balance_offset=2,
+            assign_as_block=False,
+            constraint=[
+                #{
+                #    "name": "max_shifts_per_week",
+                #    "options": 3,
+                #}
+            ],
+        ),
+        SingleShift(
+            sites=("rota a",),
+            name="weekend",
+            workers_required=1,
+            length=12.5,
+            days=days[5:],
+            balance_offset=2,
+            constraint=[
+                #{"name": "pre", "options": 2},
+                #{"name": "post", "options": 1},
+            ],
+            display_char="a",
+            #force_assign_with=["oncall"]
+        ),
+        SingleShift(
+            sites=("rota b",),
+            name="weekend b",
+            workers_required=1,
+            length=12.5,
+            days=days[5:],
+            balance_offset=2,
+            constraint=[
+                {"name": "pre", "options": 1},
+                {"name": "post", "options": 1},
+            ],
+            display_char="b",
+        ),
+        SingleShift(
+            sites=(
+                "rota a",
+                "rota b",
+            ),
+            name="twilight",
+            workers_required=1,
+            length=12.5,
+            days=days[:5],
+            balance_offset=5,
+            #constraint=[
+            #    {"name": "pre", "options": 1},
+            #    {"name": "post", "options": 1},
+            #],
+        ),
+    )
+    #Rota.allow_shifts_together_for_all_workers(
+    #        "oncall",
+    #        "twilight",
+    #)
+    #Rota.allow_shifts_together_for_all_workers(
+    #        "oncall",
+    #        "twilight",
+    #)
+
+    # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"])
+
+    load_leave = True
+    Rota.build_shifts()
+    # Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names())
+
+    for n in range(1, 12):
+        w = Worker(name=f"A{n:02d}", site="rota a", grade=1)
+        #w.prefer_multi_shift_together = 10
+        #w.allow_shifts_together("oncall", "twilight")
+        #w.allow_shifts_together("oncall", "weekend")
+
+        #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
+
+        w.add_force_assign_with("twilight", "oncall")
+        w.add_force_assign_with("weekend", "oncall")
+        w.set_max_shifts_per_week("twilight", 1)
+
+        if n <= 2:
+            pass
+            w.add_hard_day_dependency("Fri", "Sat", "Sun")
+        else:
+            w.add_hard_day_dependency("Fri", "Sun")
+            #w.set_max_shifts_per_week("weekend", 1)
+        #    #w.add_hard_day_dependency("Sat", "Sun")
+
+        print(w)
+
+        Rota.add_worker(w)
+
+    w13 = Worker(name=f"A13", site="rota a", grade=1, start_date="2025-10-01")
+    #w.prefer_multi_shift_together = 10
+    #w.allow_shifts_together("oncall", "twilight")
+    #w.allow_shifts_together("oncall", "weekend")
+
+    #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
+
+    w13.add_force_assign_with("twilight", "oncall")
+    w13.add_force_assign_with("weekend", "oncall")
+    #w.set_max_shifts_per_week("twilight", 1)
+
+    Rota.add_worker(w13)
+
+    w13 = Worker(name=f"A14", site="rota a", grade=1, start_date="2025-12-01")
+    #w.prefer_multi_shift_together = 10
+    #w.allow_shifts_together("oncall", "twilight")
+    #w.allow_shifts_together("oncall", "weekend")
+
+    #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
+
+    w13.add_force_assign_with("twilight", "oncall")
+    w13.add_force_assign_with("weekend", "oncall")
+    #w.set_max_shifts_per_week("twilight", 1)
+
+    Rota.add_worker(w13)
+
+    for w in range(14, 24):
+        w = Worker(
+            name=f"B{w}",
+            site="rota b",
+            grade=1,
+            nwds=[
+                NonWorkingDays(
+                    day="Fri",
+                )
+            ],
+        )
+        #w.add_hard_day_exclusion("Sat", "Sun")
+        w.set_max_shifts_per_week("twilight", 2)
+        w.set_max_shifts_per_week("weekend b", 1)
+
+        print(w)
+
+        Rota.add_worker(w)
+
+    # Rota.build_workers()
+    # Rota.build_model()
+
+    solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
+    # solver_options = {"seconds": time_to_run, "threads": 10}
+
+    # start_time = time.time()
+    Rota.build_and_solve(solver_options, export=True, solve=solve, solver="appsi_highs")
+
+    # Rota.solve_shifts_by_block(solver_options, block_length=13)
+    # Rota.solve_shifts_individually(solver_options)
+    # Rota.solve_model(options=solver_options)
+    # end_time = time.time()
+
+    # print(f"Time taken {end_time-start_time}")
+
+    # print(Rota.get_worker_details())
+
+    # optimizer = SolverFactory('cbc')
+    # result = optimizer.solve(prob,tee=True)
+    # result.Solver.Status = SolverStatus.warning
+    # prob.solutions.load_from(result)
+
+    # ResultsHolder = RotaResults(Rota)
+
+    # worker_timetable_brief = ResultsHolder.get_worker_timetable_brief(
+    #    show_prefs=False, show_unavailable=False
+    # )
+    #
+    # print(worker_timetable_brief)
+
+    # Rota.export_rota_to_html("proc_rota")
+    # Rota.export_rota_to_csv("rota")
+    subprocess.run(
+        ["scp", Rota.exported_rota_file, "ross@46.101.13.46:proc/proc-rota/output/"]
+    )
+
+    if suspend_on_finish:
+        os.system("systemctl suspend")
+
+
+if __name__ == "__main__":
+    app()

leave.py

@@ -27,7 +27,8 @@ def load_leave(Rota):
         }
 
         if live_rota:
-            download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSRx9VWXSlRubPyA0RhiI-Oqf5eHNYYEc6rFzlraDbR5_8qqr5g13-4uV-gn4u-TjZxiSMv1fBUaESq/pub?gid=814517272&single=true&output=csv", headers=headers)
+            download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSfcQi5Qs__A8hE2CVvUA_6ULNJALrtZOBHPUZ3xyrxPif9obtQF2IqhioT_4nebjHV1Ac5iqhNtuq4/pub?gid=2024511103&single=true&output=csv", headers=headers)
+            #download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSRx9VWXSlRubPyA0RhiI-Oqf5eHNYYEc6rFzlraDbR5_8qqr5g13-4uV-gn4u-TjZxiSMv1fBUaESq/pub?gid=814517272&single=true&output=csv", headers=headers)
             decoded_content = download.content.decode('utf-8')
 
             reader = csv.reader(decoded_content.splitlines(), delimiter=',')
@@ -56,6 +57,7 @@ def load_leave(Rota):
                 names = r
                 order = names
                 for name in names:
+                    name = name.strip()
 
                     workers[name] = {}
                     workers[name]["leave"] = []
@@ -67,13 +69,16 @@ def load_leave(Rota):
                     workers[name]["start_date"] = None
                     workers[name]["pair"] = None
 
+                n = n + 1
+                continue
+
             for i in range(len(r)):
                 try:
-                    a = order[i]
+                    a = order[i].strip()
                 except IndexError:
                     continue
 
-                worker = workers[order[i]]
+                worker = workers[a]
 
                 lower_item = r[i].lower()
                 if lower_item == "derriford":
@@ -82,16 +87,18 @@ def load_leave(Rota):
                 if lower_item == "derriford twilights":
                     lower_item = "plymouth_twilights"
 
-                if "PROC site" in row_title or row_title == "PROC nights from":
+                if "PROC site" in row_title or row_title == "PROC nights from" or "proc nights" in row_title.lower():
                     worker["site_pref"] = lower_item
 
                 elif row_title in ("Rotation", "Site", "Placement", "Placement location"):
                     worker["site"] = lower_item
 
-                elif row_title in ("Grade", "Grade (ST)"):
-                    worker["grade"] = lower_item
+                elif row_title in ("Grade", "Grade (ST)") or "Year of training" in row_title:
+                    worker["grade"] = lower_item.lstrip("ST")
 
-                elif row_title in ("%FTE", "%FTE on-call", "FTE", "FTE on-call"):
+                elif row_title in ("%FTE", "%FTE on-call", "FTE", "FTE on-call") or "on-call commitment" in row_title:
+                    if lower_item.endswith("%"):
+                        lower_item = float(lower_item.rstrip("%")) / 100
                     worker["fte"] = lower_item
 
                 elif "NWD" in row_title:
@@ -102,7 +109,10 @@ def load_leave(Rota):
 
                 elif row_title in ("End Date", "CCT date"):
                     if lower_item:
-                        date = datetime.strptime(lower_item, "%d/%m/%y").date()
+                        try:
+                            date = datetime.strptime(lower_item, "%d/%m/%y").date()
+                        except ValueError:
+                            date = datetime.strptime(lower_item, "%d/%m/%Y").date()
                         worker["end_date"] = date
                     else:
                         worker["end_date"] = None
@@ -121,7 +131,7 @@ def load_leave(Rota):
                     else:
                         worker["start_date"] = None
 
-                elif "OOP" in row_title:
+                elif "OOP" in row_title or "out of programme" in row_title.lower():
                     worker["oop"] = lower_item
 
                 elif row_title == "Group":
@@ -178,6 +188,8 @@ def load_leave(Rota):
 
             n = n + 1
 
+        
+
         return workers
 
 

output/tests/timetable.css

@@ -0,0 +1 @@
+/home/ross/proc_rota/output/timetable.css

output/tests/timetable.js

@@ -0,0 +1 @@
+/home/ross/proc_rota/output/timetable.js

output/timetable.css

@@ -321,4 +321,8 @@ table.transposed th.bank-holiday {
 
 .locum-shift {
     background-color: lightgreen;
+}
+
+.multi-shift {
+    line-height: 1;
 }

output/timetable.js

@@ -43,7 +43,16 @@ function generateExtra() {
             shifts = []
 
             shift_tds.each((n, td) => {
-                shifts.push($(td).attr("data-shift"));
+                let data_shift = $(td).attr("data-shift");
+                if (data_shift && data_shift !== "") {
+                    // Split by comma for multiple shifts per day
+                    let split_shifts = data_shift.split(",");
+                    split_shifts.forEach(s => {
+                        if (s && s.trim() !== "") {
+                            shifts.push(s.trim());
+                        }
+                    });
+                }
             })
 
             shifts_unique = new Set(shifts);
@@ -286,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];
@@ -362,11 +374,12 @@ oshifts.forEach((shift) => {
         $("table#main-table th.date").each((n, th) => {
             row = $("<tr>")
             row.append(`<td>${th.dataset.date}</td>`);
-            $(`table#main-table td[data-date='${th.dataset.date}'][data-shift='${evt.target.dataset.shift}']`).each((n, td) => {
-
-                //row.append(`<td>${}</td>`)
-                row.append($(td).closest("tr").children("td:first").clone())
-
+            $(`table#main-table td[data-date='${th.dataset.date}']`).filter(function() {
+                let shifts = $(this).attr("data-shift");
+                if (!shifts) return false;
+                return shifts.split(",").map(s => s.trim()).includes(evt.target.dataset.shift);
+            }).each((n, td) => {
+                row.append($(td).closest("tr").children("td:first").clone());
             });
 
             table.append(row)
@@ -416,8 +429,9 @@ $(".table-div + pre").each((n, pre) => {
 
 })
 
+
 $("td").hover((e) => {
-    //start hover
+    // start hover
     worker_td = $(e.target).closest('tr').find('td:first')[0]
     pair = worker_td.dataset.pair;
     index = $(e.target).index() - 1;
@@ -426,24 +440,36 @@ $("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) {
-        $(`td[data-shift=${e.target.dataset.shift}]`).addClass("shift-highlight")
 
+    // 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);
 
-        $(`td[data-remote-site='${e.target.dataset.shiftRemoteSite}']`).closest('tr').find(`td[data-shift=${e.target.dataset.shift}]`).addClass("remote-site-match")
-        // if (e.target.dataset.shiftRemoteSite == worker_td.dataset.remoteSite) {
-        //     $(`td[data-shift-remote-site=${e.target.dataset.shiftRemoteSite}]`).addClass("remote-site-match")
-        // }
+    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");
+            if (!shifts) return false;
+            let cell_shifts = shifts.split(",").map(s => s.trim());
+            return hovered_shifts.some(s => cell_shifts.includes(s));
+        }).addClass("shift-highlight");
     }
 
 }, (e) => {
-
-    //end hover
+    // 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(() => {

rota/workers.py

@@ -111,9 +111,23 @@ 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
+    hard_day_dependencies: set[frozenset[str]] = set()  # e.g. {frozenset({"Mon", "Tue"}), frozenset({"Fri", "Sun"})}
+    hard_day_exclusions: list[tuple[str, str]] = []  # e.g. [("Fri", "Sun")]
+    force_assign_with: dict[str, list[str]] = {}  # e.g. {"oncall": ["weekend"]}
+    max_shifts_per_week_by_shift_name: dict[str, int] = {}  # e.g. {"night": 2, "a": 3}
+
+
+    assign_as_block_preferences: dict[str, float] = {}  # shift_name -> weight (positive = prefer block)
 
     shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist
 
+    weekend_shift_target_number: int = 0
+
+
+
     model_config = ConfigDict(
         extra="allow",
     )
@@ -223,7 +237,7 @@ class Worker(BaseModel):
                 end_oop_date = datetime.datetime.strptime(end_oop, "%d/%m/%y").date()
 
             if start_oop_date >= end_oop_date:
-                raise ValueError("End OOP date must be after start date")
+                raise ValueError(f"End OOP date must be after start date [{self.name} - Start date: {self.start_date} / End date: {self.end_date}]")
 
             # ignore oops if they finish before the rota (or worker) start date
             if end_oop_date > self.calculated_start_date:
@@ -389,3 +403,31 @@ 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))
+
+    def add_hard_day_dependency(self, *days: str):
+        if not hasattr(self, "hard_day_dependencies"):
+            self.hard_day_dependencies = set()
+        self.hard_day_dependencies.add(frozenset(days))
+
+    def add_hard_day_exclusion(self, day1: str, day2: str):
+        if not hasattr(self, "hard_day_exclusions"):
+            self.hard_day_exclusions = []
+        self.hard_day_exclusions.append((day1, day2))
+
+    def add_force_assign_with(self, shift: str, with_shifts: list[str] | str):
+        if not hasattr(self, "force_assign_with"):
+            self.force_assign_with = {}
+        if isinstance(with_shifts, str):
+            with_shifts = [with_shifts]
+        self.force_assign_with[shift] = with_shifts
+
+    def set_max_shifts_per_week(self, shift_name: str, max_per_week: int):
+        """Helper to set the maximum number of shifts of a given type per week for this worker."""
+        if not hasattr(self, "max_shifts_per_week_by_shift_name"):
+            self.max_shifts_per_week_by_shift_name = {}
+        self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week

test/rota_constraints/test_summed_shift_grade_constraints.py

@@ -0,0 +1,94 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days, MinSummedGradeByShiftsPerDayConstraint
+
+from rota.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'"

test/test_balance_shift_sites.py

@@ -0,0 +1,139 @@
+import pytest
+import datetime
+from rota.shifts import RotaBuilder, SingleShift, days
+from rota.workers import Worker
+
+@pytest.fixture
+def demo_rota_balance_sites():
+    weeks_to_rota = 8
+    start_date = datetime.date(2022, 3, 7)
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    worker1 = Worker(name="worker1", site="group1", grade=1)
+    worker2 = Worker(name="worker2", site="group1", grade=1)
+    worker3 = Worker(name="worker3", site="group2", grade=1)
+    worker4 = Worker(name="worker4", site="group2", grade=1)
+    worker5 = Worker(name="worker5", site="group3", grade=1)
+    worker6 = Worker(name="worker6", site="group3", grade=1)
+    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6))
+    return Rota
+
+def test_balance_blocks_across_2_groups(demo_rota_balance_sites):
+    Rota = demo_rota_balance_sites
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="weekday_night",
+            length=12.5,
+            days=days[:4],
+            balance_offset=40,
+            workers_required=2,
+            force_as_block=True,
+            constraint=[{"name": "balance_across_groups"}],
+        ),
+    )
+    Rota.constraint_options["balance_nights_across_sites"] = False
+    Rota.constraint_options["balance_nights"] = False
+    Rota.constraint_options["constrain_time_off_after_nights"] = False
+
+    # Worker 5 and 6 don't have any valid shifts so this should fail
+    #with pytest.raises(WarningTermination):
+    #    Rota.build_and_solve(options={"ratio": 0.0})
+
+    # Unless we remove the warning
+    Rota.terminate_on_warning.remove("Worker/no valid shifts")
+    Rota.build_and_solve(options={"ratio": 0.0})
+
+    assert len(Rota.get_warnings("Worker/no valid shifts")) == 2
+
+    Rota.export_rota_to_html("test5")
+
+    group_workers = Rota.get_workers_by_group()
+    for group in group_workers:
+        shift_patterns = []
+        for w in group_workers[group]:
+            shift_patterns.append(Rota.get_worker_shift_list(w))
+
+        zipped_lists = list(zip(*shift_patterns))
+
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") <= 1
+
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+
+def test_balance_blocks_across_2_groups_unbalanced(demo_rota_balance_sites):
+    Rota = demo_rota_balance_sites
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="weekday_night",
+            length=12.5,
+            days=days[:4],
+            balance_offset=40,
+            workers_required=3,
+            force_as_block=True,
+            constraint=[{"name": "balance_across_groups"}],
+        ),
+    )
+    Rota.constraint_options["balance_nights_across_sites"] = False
+    Rota.constraint_options["balance_nights"] = False
+    Rota.constraint_options["constrain_time_off_after_nights"] = False
+
+    try:
+        Rota.terminate_on_warning.remove("Worker/no valid shifts")
+    except ValueError:
+        # Seems to happen with sequential test running
+        pass
+    Rota.build_and_solve(options={"ratio": 0.1})
+
+    group_workers = Rota.get_workers_by_group()
+    for group in group_workers:
+        shift_patterns = []
+        for w in group_workers[group]:
+            shift_patterns.append(Rota.get_worker_shift_list(w))
+
+        zipped_lists = list(zip(*shift_patterns))
+
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") <= 2
+
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+
+def test_balance_blocks_across_3_groups(demo_rota_balance_sites):
+    Rota = demo_rota_balance_sites
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days[:4],
+            balance_offset=40,
+            workers_required=3,
+            force_as_block=True,
+            constraint=[{"name": "balance_across_groups"}],
+        ),
+    )
+    Rota.constraint_options["balance_nights_across_sites"] = False
+    Rota.constraint_options["balance_nights"] = False
+    Rota.constraint_options["constrain_time_off_after_nights"] = False
+    Rota.build_and_solve(options={"ratio": 0.1})
+
+    group_workers = Rota.get_workers_by_group()
+    for group in group_workers:
+        shift_patterns = []
+        for w in group_workers[group]:
+            shift_patterns.append(Rota.get_worker_shift_list(w))
+
+        zipped_lists = list(zip(*shift_patterns))
+
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") <= 1
+
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"

test/test_general_balancing.py

@@ -1,284 +1,227 @@
 import datetime
 import pytest
 from pytest import approx
-from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
-
+from rota.shifts import 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
 
+def test_basic_balancing():
+    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]
+        assert worker_shifts["a"] in (approx(22), approx(23), approx(24), approx(25))
+        assert worker_shifts["b"] in (approx(22), approx(23), approx(24), approx(25))
 
-class TestBalancing:
-    def test_basic_balancing(self):
-        Rota = generate_basic_rota()
+def test_weighted_shift_balancing():
+    Rota = generate_basic_rota(20)
+    Rota.add_shifts(
+        SingleShift(
+            sites=["group1", "group2"],
+            name="a",
+            length=12.5,
+            days=days,
+            balance_weighting=10,
+            workers_required=1,
+            force_as_block=False,
+            constraint=[{"name": "pre","options": "2"}, {"name": "post","options": "2"}],
+        ),
+        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.0001})
+    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]
+        assert worker_shifts["a"] in (approx(46), approx(47), approx(48))
 
-        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,
-            ),
-        )
+def test_weighted_shift_balancing2():
+    Rota = generate_basic_rota(23)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="a",
+            length=12.5,
+            days=days,
+            workers_required=1,
+            force_as_block=False,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="b",
+            length=12.5,
+            days=days,
+            balance_weighting=8,
+            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.0001})
+    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]
+        assert worker_shifts["b"] in (approx(53), approx(54), approx(55))
 
-        Rota.build_and_solve(options={"ratio": 0.001})
-        Rota.export_rota_to_html("basic_balancing")
+def test_weighted_shift_balancing3():
+    Rota = generate_basic_rota(23)
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="a",
+            length=12.5,
+            days=days,
+            workers_required=1,
+            force_as_block=False,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="b",
+            length=12.5,
+            days=days,
+            balance_weighting=1,
+            workers_required=1,
+            force_as_block=False,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="c",
+            length=12.5,
+            days=days[0],
+            balance_weighting=8,
+            workers_required=1,
+            force_as_block=False,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0001})
+    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]
+        assert worker_shifts["c"] in (approx(7), approx(8))
 
-        shift_summary = Rota.get_shift_summary_dict()
-        for worker_name in shift_summary:
-            worker_shifts = shift_summary[worker_name]
+def test_weighted_shift_balancing4():
+    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,
+        ),
+        SingleShift(
+            sites=("group2", "group3"),
+            name="b",
+            length=12.5,
+            days=days,
+            balance_offset=99,
+            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]
+        assert worker_shifts["a"] in (0, 46, 47, 48)
 
-            #worker = Rota.get_worker_by_name(worker_name)
-            assert worker_shifts["a"] in (approx(22), approx(23), approx(24), approx(25))
-            assert worker_shifts["b"] in (approx(22), approx(23), approx(24), approx(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=10,
-                workers_required=1,
-                force_as_block=False,
-                constraint=[{"name": "pre","options": "2"}, {"name": "post","options": "2"}],
-            ),
-            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.0001})
-        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 (approx(46), approx(47), approx(48))
-            #assert worker_shifts["b"] 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=8,
-                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.0001})
-        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 (approx(53), approx(54), approx(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=1,
-                workers_required=1,
-                force_as_block=False,
-            ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="c",
-                length=12.5,
-                days=days[0],
-                balance_weighting=8,
-                workers_required=1,
-                force_as_block=False,
-            ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.0001})
-        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 (approx(7), approx(8))
-            #assert worker_shifts["b"] in (53, 54)
-
-    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)
+def test_weighted_shift_balancing5():
+    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,
+            workers_required=2,
+            force_as_block=False,
+        ),
+        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]
+        assert worker_shifts["b"] in (0, 35)

test/test_limit_constraints.py

@@ -0,0 +1,325 @@
+import pytest
+import datetime
+from rota.shifts import RotaBuilder, SingleShift, days
+from rota.workers import Worker
+
+@pytest.fixture
+def limit_constraint_rota():
+    weeks_to_rota = 8
+    start_date = datetime.date(2022, 3, 7)
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    worker1 = Worker(name="worker1", site="group1", grade=2, remote_site="group1")
+    worker2 = Worker(name="worker2", site="group1", grade=2, remote_site="group1")
+    worker3 = Worker(name="worker3", site="group1", grade=2, remote_site="group1")
+    worker4 = Worker(name="worker4", site="group3", grade=3, remote_site="group2")
+    worker5 = Worker(name="worker5", site="group3", grade=3, remote_site="group2")
+    worker6 = Worker(name="worker6", site="group3", grade=3, remote_site="group2")
+    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6))
+    Rota.shifts = []
+    return Rota
+
+def test_constraint_limit_grades(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=60,
+            workers_required=4,
+            force_as_block=True,
+            constraint=[{"name": "limit_grade_number", "options": {2: 1}}],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.01})
+    Rota.export_rota_to_html("constraint_limit_grades", folder="tests")
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+    grade_workers = Rota.get_workers_by_grade()
+    for grade in grade_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
+        zipped_lists = list(zip(*shift_patterns))
+        limit = 1 if grade == 2 else 3
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") == limit
+
+def test_constraint_limit_grades2(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraint=[{"name": "limit_grade_number", "options": {3: 1}}],
+        ),
+    )
+    Rota.constraint_options["balance_shifts_quadratic"] = True
+    Rota.build_and_solve(options={"ratio": 0.1})
+    Rota.export_rota_to_html("test9", folder="tests")
+    grade_workers = Rota.get_workers_by_grade()
+    for grade in grade_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
+        zipped_lists = list(zip(*shift_patterns))
+        limit = 1 if grade == 3 else 3
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") == limit
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+
+def test_constraint_limit_grades3(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.workers = []
+    worker1 = Worker(name="worker1", site="group1", grade=2, remote_site="group1")
+    worker2 = Worker(name="worker2", site="group1", grade=2, remote_site="group1")
+    worker3 = Worker(name="worker3", site="group1", grade=2, remote_site="group1")
+    worker4 = Worker(name="worker4", site="group3", grade=3, remote_site="group2")
+    worker5 = Worker(name="worker5", site="group3", grade=3, remote_site="group2")
+    worker6 = Worker(name="worker6", site="group3", grade=3, remote_site="group2")
+    worker7 = Worker(name="worker7", site="group3", grade=2, remote_site="group2")
+    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6, worker7))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=60,
+            workers_required=4,
+            force_as_block=True,
+            constraint=[{"name": "limit_grade_number", "options": {2: 3, 3: 1}}],
+        ),
+    )
+    Rota.constraint_options["balance_shifts_quadratic"] = True
+    Rota.build_and_solve(options={"ratio": 0.01})
+    Rota.export_rota_to_html("test9")
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+    grade_workers = Rota.get_workers_by_grade()
+    for grade in grade_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
+        zipped_lists = list(zip(*shift_patterns))
+        limit = 3 if grade == 2 else 1
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") == limit
+
+def test_constraint_limit_grades4(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraint=[{"name": "limit_grade_number", "options": {2: 4, 3: 0}}],
+        ),
+    )
+    Rota.constraint_options["balance_shifts_quadratic"] = True
+    Rota.build_and_solve(options={"ratio": 0.1})
+    assert Rota.results.solver.termination_condition == "infeasible"
+
+def test_constraint_minimum_grades(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=1,
+            force_as_block=True,
+            constraint=[{"name": "minimum_grade_number", "options": (3, 1)}],
+        ),
+    )
+    Rota.constraint_options["balance_shifts_quadratic"] = True
+    Rota.build_and_solve(options={"ratio": 0.1})
+    Rota.export_rota_to_html("test9")
+    assert Rota.results.solver.status == "ok"
+    grade_workers = Rota.get_workers_by_grade()
+    for grade in grade_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
+        zipped_lists = list(zip(*shift_patterns))
+        limit = 1 if grade > 2 else 0
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") == limit
+
+def test_constraint_minimum_grades2(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraint=[{"name": "minimum_grade_number", "options": (3, 3)}],
+        ),
+    )
+    Rota.constraint_options["balance_shifts_quadratic"] = True
+    Rota.build_and_solve(options={"ratio": 0.1})
+    Rota.export_rota_to_html("test9")
+    grade_workers = Rota.get_workers_by_grade()
+    for grade in grade_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
+        zipped_lists = list(zip(*shift_patterns))
+        limit = 3 if grade > 2 else 1
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("weekday_night") == limit
+    assert Rota.results.solver.status == "ok"
+
+def test_constraint_minimum_grades3(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraint=[{"name": "minimum_grade_number", "options": (3, 0)}],
+        ),
+    )
+    Rota.constraint_options["balance_shifts_quadratic"] = True
+    Rota.build_and_solve(options={"ratio": 0.1})
+    Rota.export_rota_to_html("test9")
+    assert Rota.results.solver.status == "ok"
+
+def test_constraint_minimum_grades_no_valid_worker(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="weekday_night",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraint=[{"name": "minimum_grade_number", "options": (4, 1)}],
+        ),
+    )
+    Rota.constraint_options["balance_shifts_quadratic"] = True
+    with pytest.raises(ValueError):
+        Rota.build_and_solve(options={"ratio": 0.1})
+
+def test_constraint_require_remote_site_presence_week(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="night",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=2,
+            force_as_block=True,
+            constraint=[
+                {
+                    "name": "require_remote_site_presence_week",
+                    "options": ("group1", 2),
+                }
+            ],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("remote1")
+    group_workers = Rota.get_workers_by_group()
+    for group in group_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in group_workers[group]]
+        zipped_lists = list(zip(*shift_patterns))
+        limit = 2 if group == "group1" else 0
+        for day_shifts in zipped_lists:
+            assert day_shifts.count("night") == limit
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+
+def test_constraint_require_remote_site_presence_week2(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.constraint_options["minimise_shift_diffs"] = False
+    Rota.constraint_options["balance_shifts_quadratic"] = False
+    Rota.constraint_options["balance_shift"] = False
+    Rota.constraint_options["balance_nights_across_sites"] = True
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="night",
+            length=12.5,
+            days=days,
+            balance_offset=20,
+            workers_required=3,
+            force_as_block=True,
+            constraint=[
+                {
+                    "name": "require_remote_site_presence_week",
+                    "options": ("group1", 2),
+                }
+            ],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("remote2")
+    group_workers = Rota.get_workers_by_remote_group()
+    for group in group_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in group_workers[group]]
+        zipped_lists = list(zip(*shift_patterns))
+        if group == "group1":
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("night") >= 2
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+
+def test_constraint_require_remote_site_presence_week3(limit_constraint_rota):
+    Rota = limit_constraint_rota
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2", "group3"),
+            name="night",
+            length=12.5,
+            days=days,
+            balance_offset=20,
+            workers_required=3,
+            force_as_block=True,
+            constraint=[
+                {
+                    "name": "require_remote_site_presence_week",
+                    "options": ("group2", 1),
+                }
+            ],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("remote3")
+    group_workers = Rota.get_workers_by_remote_group()
+    for group in group_workers:
+        shift_patterns = [Rota.get_worker_shift_list(w) for w in group_workers[group]]
+        zipped_lists = list(zip(*shift_patterns))
+        if group == "group2":
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("night") >= 1
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"

test/test_multiple_shifts_per_day.py

@@ -0,0 +1,276 @@
+
+import pytest
+import datetime
+from rota.shifts import RotaBuilder, SingleShift, days
+from rota.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"

test/test_night_unavailable.py

@@ -0,0 +1,146 @@
+import pytest
+import datetime
+from rota.shifts import RotaBuilder, SingleShift, days
+from rota.workers import Worker
+
+@pytest.fixture
+def demo_rota_night_unavailable():
+    weeks_to_rota = 10
+    start_date = datetime.date(2022, 3, 7)
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    worker1 = Worker(
+        name="worker1",
+        site="group1",
+        grade=1,
+        not_available_to_work=[
+            {
+                "date": datetime.datetime.strptime("15/03/22", "%d/%m/%y").date(),
+                "reason": "****",
+            },
+        ],
+    )
+    worker2 = Worker(
+        name="worker2",
+        site="group1",
+        grade=1,
+        not_available_to_work=[
+            {
+                "date": datetime.datetime.strptime("14/03/22", "%d/%m/%y").date(),
+                "reason": "****",
+            }
+        ],
+    )
+    Rota.add_workers((worker1, worker2))
+    return Rota
+
+def test_basic_assignment(demo_rota_night_unavailable):
+    Rota = demo_rota_night_unavailable
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="night_weekend",
+            length=12.5,
+            days=days[5:],
+            constraint=[{"name": "night"}],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("testnight")
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+    summary = Rota.get_shift_summary_dict()
+    for worker in summary:
+        assert summary[worker]["night_weekend"] == 10
+
+def test_assign_night_prior_to_unavailablity(demo_rota_night_unavailable):
+    Rota = demo_rota_night_unavailable
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="night_weekend",
+            length=12.5,
+            days=days[5:],
+            workers_required=2,
+            constraint=[{"name": "night"}],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("testnight_unavail")
+    assert Rota.results.solver.status in ("warning", "error")
+    assert Rota.results.solver.termination_condition == "infeasible"
+
+def test_assign_non_night_prior_to_unavailablity(demo_rota_night_unavailable):
+    Rota = demo_rota_night_unavailable
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="night_weekend",
+            length=12.5,
+            days=days[5:],
+            workers_required=2,
+            #constraint=[{"name": "night"}],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("testnight_unavail")
+    assert Rota.results.solver.status == "ok"
+
+def test_assign_prior_to_unavailablity_non_night(demo_rota_night_unavailable):
+    Rota = demo_rota_night_unavailable
+    Rota.shifts = []
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="night_weekend",
+            length=12.5,
+            days=days[5:],
+            workers_required=2,
+            constraint=[],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("testnight")
+    assert Rota.results.solver.status == "ok"
+
+def test_assign_split(demo_rota_night_unavailable):
+    Rota = demo_rota_night_unavailable
+    Rota.shifts = []
+    worker3 = Worker(
+        name="worker3",
+        site="group1",
+        grade=1,
+        not_available_to_work=(
+            {"date": datetime.datetime.strptime("13/03/22", "%d/%m/%y").date()},
+        ),
+    )
+    worker4 = Worker(
+        name="worker4",
+        site="group1",
+        grade=1,
+        not_available_to_work=(
+            {"date": datetime.datetime.strptime("12/03/22", "%d/%m/%y").date()},
+        ),
+    )
+    Rota.add_workers((worker3, worker4))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="night_weekend",
+            length=12.5,
+            days=days[5:],
+            workers_required=3,
+            constraint=[],
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("testnight")
+    assert Rota.results.solver.status == "ok"
+    summary = Rota.get_shift_summary_dict()
+    for worker in summary:
+        assert summary[worker]["night_weekend"] == 15

test/test_nights2.py

@@ -1,512 +1,371 @@
 import datetime
 import pytest
 from rota.shifts import InvalidShift, 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,
     )
-
-    # Add a few workers
     Rota.add_workers(
         [
             Worker(name="worker1", site="group1", grade=1),
             Worker(name="worker2", site="group1", grade=1),
         ]
     )
-
     return Rota
 
-
 def date_generator(from_date, days):
     n = 0
     while True:
         yield from_date
-
         n = n + 1
         if n >= days:
             break
         from_date = from_date + datetime.timedelta(days=1)
 
+def test_nights():
+    Rota = generate_basic_rota()
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}, {"name": "pre", "options": 2}],
+            workers_required=2,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="b",
+            length=12.5,
+            days=days[:3],
+            force_as_block=False,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night", folder="tests")
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
 
-class TestNightShifts:
-    def test_nights(self):
-        Rota = generate_basic_rota()
+def test_nights_pre3():
+    Rota = generate_basic_rota()
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+            workers_required=2,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="b",
+            length=12.5,
+            days=days[:2],
+            force_as_block=False,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night_pre3", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}, {"name": "pre", "options": 2}],
-                workers_required=2,
-            ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="b",
-                length=12.5,
-                days=days[:3],
-                force_as_block=False,
-                # constraint=[{"name": "night"}],
-            ),
-        )
+def test_nights_pre3_2():
+    Rota = generate_basic_rota()
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+            workers_required=1,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="b",
+            length=12.5,
+            days=days[:3],
+            force_as_block=False,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night_pre3", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
+def test_nights_fail():
+    Rota = generate_basic_rota()
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+            workers_required=2,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="b",
+            length=12.5,
+            days=days[:3],
+            force_as_block=False,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    assert Rota.results.solver.status in ("warning", "error")
 
-        assert Rota.results.solver.status == "ok"
-        assert Rota.results.solver.termination_condition == "optimal"
+def test_nights2():
+    Rota = generate_basic_rota()
+    Rota.constraint_options["max_night_frequency"] = 0
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=2,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="m",
+            length=12.5,
+            days=days[:5],
+            force_as_block=True,
+            workers_required=2,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-    def test_nights_pre3(self):
-        Rota = generate_basic_rota()
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-                workers_required=2,
-            ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="b",
-                length=12.5,
-                days=days[:2],
-                force_as_block=False,
-                # constraint=[{"name": "night"}],
-            ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night_pre3")
-
-        assert Rota.results.solver.status == "ok"
-
-    def test_nights_pre3_2(self):
-        Rota = generate_basic_rota()
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-                workers_required=1,
-            ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="b",
-                length=12.5,
-                days=days[:3],
-                force_as_block=False,
-                # constraint=[{"name": "night"}],
-            ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night_pre3")
-
-        assert Rota.results.solver.status == "ok"
-
-    def test_nights_fail(self):
-        Rota = generate_basic_rota()
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-                workers_required=2,
-            ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="b",
-                length=12.5,
-                days=days[:3],
-                force_as_block=False,
-                # constraint=[{"name": "night"}],
-            ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-
-        assert Rota.results.solver.status in ("warning", "error")
-
-    def test_nights2(self):
-        Rota = generate_basic_rota()
-
-        Rota.constraint_options["max_night_frequency"] = 0
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}],
-                workers_required=2,
-            ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="m",
-                length=12.5,
-                days=days[:5],
-                force_as_block=True,
-                # constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-                # constraint=[{"name": "night"}],
-                workers_required=2,
-            ),
-            # SingleShift(
-            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
-            #    force_as_block=False,
-            #    #constraint=[{"name": "night"}],
-            # ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status == "ok"
-
-    def test_nights_pre_wrap_around(self):
-        Rota = generate_basic_rota(weeks_to_rota=8)
-
-        Rota.constraint_options["max_night_frequency"] = 0
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[:2],
-                force_as_block=True,
-                constraint=[
-                    {"name": "night"},
-                    {"name": "pre", "options": 5},
-                    {"name": "post", "options": 5},
+def test_nights_pre_wrap_around():
+    Rota = generate_basic_rota(weeks_to_rota=8)
+    Rota.constraint_options["max_night_frequency"] = 0
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[:2],
+            force_as_block=True,
+            constraint=[
+                {"name": "night"},
+                {"name": "pre", "options": 5},
+                {"name": "post", "options": 5},
+            ],
+            workers_required=3,
+        ),
+    )
+    Rota.add_workers(
+        [
+            Worker(
+                name="worker3",
+                site="group1",
+                grade=1,
+                not_available_to_work=[
+                    {
+                        "date": Rota.rota_end_date - datetime.timedelta(days=1),
+                        "reason": "no wanna",
+                    },
                 ],
-                workers_required=3,
             ),
-            # SingleShift(
-            #    sites=("group1", "group2"), name="m", length= 12.5, days=days[:5],
-            #    force_as_block=True,
-            #    #constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-            #    #constraint=[{"name": "night"}],
-            #    workers_required=2,
-            # ),
-            # SingleShift(
-            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
-            #    force_as_block=False,
-            #    #constraint=[{"name": "night"}],
-            # ),
-        )
+        ]
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night pre wrap around", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-        Rota.add_workers(
-            [
-                Worker(
-                    name="worker3",
-                    site="group1",
-                    grade=1,
-                    not_available_to_work=[
-                        {
-                            "date": Rota.rota_end_date - datetime.timedelta(days=1),
-                            "reason": "no wanna",
-                        },
-                        # {
-                        #    "date": Rota.rota_end_date - datetime.timedelta(days=5),
-                        #    "reason": "no wanna",
-                        # },
-                    ],
-                ),
-            ]
-        )
+def test_nights_max_frequency():
+    Rota = generate_basic_rota()
+    Rota.constraint_options["max_night_frequency"] = 1
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=2,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night max frequency", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
+def test_nights_max_frequency_fail():
+    Rota = generate_basic_rota()
+    Rota.constraint_options["max_night_frequency"] = 2
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=2,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    assert Rota.results.solver.status in ("warning", "error")
 
-        assert Rota.results.solver.status == "ok"
+def test_nights_max_frequency3():
+    Rota = generate_basic_rota()
+    Rota.constraint_options["max_night_frequency"] = 2
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=1,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="m",
+            length=12.5,
+            days=days[:5],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night max frequency 3", folder="tests")
+    assert Rota.results.solver.status in ("warning", "error")
 
-    def test_nights_max_frequency(self):
-        Rota = generate_basic_rota()
+def test_nights_max_frequency4():
+    Rota = generate_basic_rota()
+    Rota.constraint_options["max_night_frequency"] = 2
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=1,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="m",
+            length=12.5,
+            days=days[:5],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=1,
+        ),
+    )
+    Rota.add_workers(
+        [
+            Worker(name="worker3", site="group1", grade=1),
+            Worker(name="worker4", site="group1", grade=1),
+        ]
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night max frequency 4", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-        Rota.constraint_options["max_night_frequency"] = 1
+def test_nights_max_frequency_exclusions():
+    Rota = generate_basic_rota()
+    Rota.constraint_options["max_night_frequency"] = 2
+    Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(10))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=2,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night max frequency exclusions", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}],
-                workers_required=2,
+def test_nights_max_frequency_exclusions2():
+    Rota = generate_basic_rota(weeks_to_rota=8)
+    Rota.constraint_options["max_night_frequency"] = 2
+    Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(1,4))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            constraint=[{"name": "night"}],
+            workers_required=2,
+        ),
+    )
+    Rota.add_workers(
+        [
+            Worker(
+                name="worker3",
+                site="group1",
+                grade=1,
+                start_date=Rota.start_date + datetime.timedelta(weeks=4)
             ),
-            # SingleShift(
-            #    sites=("group1", "group2"),
-            #    name="m",
-            #    length=12.5,
-            #    days=days[:5],
-            #    force_as_block=True,
-            #    # constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-            #    # constraint=[{"name": "night"}],
-            #    workers_required=2,
-            # ),
-            # SingleShift(
-            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
-            #    force_as_block=False,
-            #    #constraint=[{"name": "night"}],
-            # ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status == "ok"
-
-    def test_nights_max_frequency_fail(self):
-        Rota = generate_basic_rota()
-
-        Rota.constraint_options["max_night_frequency"] = 2
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}],
-                workers_required=2,
+            Worker(
+                name="worker4",
+                site="group1",
+                grade=1,
+                start_date=Rota.start_date + datetime.timedelta(weeks=4)
             ),
-            # SingleShift(
-            #    sites=("group1", "group2"),
-            #    name="m",
-            #    length=12.5,
-            #    days=days[:5],
-            #    force_as_block=True,
-            #    # constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-            #    # constraint=[{"name": "night"}],
-            #    workers_required=2,
-            # ),
-            # SingleShift(
-            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
-            #    force_as_block=False,
-            #    #constraint=[{"name": "night"}],
-            # ),
-        )
+        ]
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night max frequency exclusions 2", folder="tests")
+    assert Rota.results.solver.status == "ok"
 
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status in ("warning", "error")
-
-    def test_nights_max_frequency3(self):
-        Rota = generate_basic_rota()
-
-        Rota.constraint_options["max_night_frequency"] = 2
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}],
-                workers_required=1,
+def test_nights_max_frequency_exclusions3():
+    Rota = generate_basic_rota(weeks_to_rota=9)
+    Rota.constraint_options["max_night_frequency"] = 2
+    Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(1,5))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="n",
+            length=12.5,
+            days=days[5:],
+            force_as_block=True,
+            balance_offset=10,
+            constraint=[{"name": "night"}],
+            workers_required=2,
+        ),
+    )
+    Rota.add_workers(
+        [
+            Worker(
+                name="worker3",
+                site="group1",
+                grade=1,
+                start_date=Rota.start_date + datetime.timedelta(weeks=4)
             ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="m",
-                length=12.5,
-                days=days[:5],
-                force_as_block=True,
-                # constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-                 constraint=[{"name": "night"}],
-                workers_required=1,
+            Worker(
+                name="worker4",
+                site="group1",
+                grade=1,
+                start_date=Rota.start_date + datetime.timedelta(weeks=4)
             ),
-            # SingleShift(
-            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
-            #    force_as_block=False,
-            #    #constraint=[{"name": "night"}],
-            # ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status in ("warning", "error")
-
-    def test_nights_max_frequency4(self):
-        Rota = generate_basic_rota()
-
-        Rota.constraint_options["max_night_frequency"] = 2
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}],
-                workers_required=1,
-            ),
-            SingleShift(
-                sites=("group1", "group2"),
-                name="m",
-                length=12.5,
-                days=days[:5],
-                force_as_block=True,
-                # constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
-                constraint=[{"name": "night"}],
-                workers_required=1,
-            ),
-            # SingleShift(
-            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
-            #    force_as_block=False,
-            #    #constraint=[{"name": "night"}],
-            # ),
-        )
-
-        Rota.add_workers(
-            [
-                Worker(
-                    name="worker3",
-                    site="group1",
-                    grade=1,
-                ),
-                Worker(
-                    name="worker4",
-                    site="group1",
-                    grade=1,
-                ),
-            ]
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status == "ok"
-
-    def test_nights_max_frequency_exclusions(self):
-        Rota = generate_basic_rota()
-
-        Rota.constraint_options["max_night_frequency"] = 2
-        Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(10))
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}],
-                workers_required=2,
-            ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status == "ok"
-
-    def test_nights_max_frequency_exclusions2(self):
-        Rota = generate_basic_rota(weeks_to_rota=8)
-
-        Rota.constraint_options["max_night_frequency"] = 2
-        Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(1,4))
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                constraint=[{"name": "night"}],
-                workers_required=2,
-            ),
-        )
-
-        Rota.add_workers(
-            [
-                Worker(
-                    name="worker3",
-                    site="group1",
-                    grade=1,
-                    start_date=Rota.start_date + datetime.timedelta(weeks=4)
-                ),
-                Worker(
-                    name="worker4",
-                    site="group1",
-                    grade=1,
-                    start_date=Rota.start_date + datetime.timedelta(weeks=4)
-                ),
-            ]
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status == "ok"
-
-    def test_nights_max_frequency_exclusions3(self):
-        # TODO: find out the issues with odd rotas and exclusions
-        Rota = generate_basic_rota(weeks_to_rota=9)
-
-        Rota.constraint_options["max_night_frequency"] = 2
-        Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(1,5))
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"),
-                name="n",
-                length=12.5,
-                days=days[5:],
-                force_as_block=True,
-                balance_offset=10,
-                constraint=[{"name": "night"}],
-                workers_required=2,
-            ),
-        )
-
-        Rota.add_workers(
-            [
-                Worker(
-                    name="worker3",
-                    site="group1",
-                    grade=1,
-                    start_date=Rota.start_date + datetime.timedelta(weeks=4)
-                ),
-                Worker(
-                    name="worker4",
-                    site="group1",
-                    grade=1,
-                    start_date=Rota.start_date + datetime.timedelta(weeks=4)
-                ),
-            ]
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
-
-        assert Rota.results.solver.status == "ok"
+        ]
+    )
+    Rota.build_and_solve(options={"ratio": 0.000})
+    Rota.export_rota_to_html("night max frequency exclusions 3", folder="tests")
+    assert Rota.results.solver.status == "ok"

test/test_nwds.py

@@ -1,374 +1,216 @@
 import pytest
 from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
-
 import datetime
-
 from rota.workers import Worker
 
-
 def weeks_from_list(lst):
     """Yield successive n-sized chunks from a lst."""
     for i in range(0, len(lst), 7):
         yield lst[i : i + 7]
 
-
-class TestDemoRota:
-    def test_nwd(self):
-        # Set up rota
-        weeks_to_rota = 10
+def setup_rota(weeks_to_rota=10, start_date=None):
+    if start_date is None:
         start_date = datetime.date(2022, 3, 7)
+    return RotaBuilder(start_date, weeks_to_rota=weeks_to_rota)
 
-        self.Rota = RotaBuilder(
-            start_date,
-            weeks_to_rota=weeks_to_rota,
-        )
+def test_nwd():
+    Rota = setup_rota()
+    worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
+    worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Fri"}])
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="d", length=12.5, days=days[:5], workers_required=1,
+        ),
+        SingleShift(
+            sites=("group1",), name="w", length=12.5, days=days[5:], workers_required=2,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.10})
+    Rota.export_rota_to_html("nwd1", folder="tests")
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+    for worker in Rota.workers:
+        shifts = Rota.get_worker_shift_list(worker)
+        shifts_string = "".join([i if i != "" else "-" for i in shifts])
+        if worker.name == "worker1":
+            assert "wwd" not in shifts_string
+        if worker.name == "worker2":
+            assert "dww" not in shifts_string
 
-        # Add a few workers
-        worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
-        worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Fri"}])
+def test_nwd_simple_fail():
+    Rota = setup_rota()
+    worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
+    worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Mon"}])
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd", folder="tests")
+    assert Rota.results.solver.status in ("warning", "error")
 
-        self.Rota.add_workers((worker1, worker2))
+def test_nwd_partial_rota():
+    Rota = setup_rota()
+    start_date = Rota.start_date
+    d = start_date + datetime.timedelta(weeks=5)
+    worker1 = Worker(
+        name="worker1", site="group1", grade=1,
+        nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
+    )
+    worker2 = Worker(
+        name="worker2", site="group1", grade=1,
+        nwds=[{"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date}],
+    )
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd", 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])
+        if worker.name == "worker1":
+            for n in range(0, 4):
+                assert shifts_string[0 + 7 * n] == "-"
+        if worker.name == "worker2":
+            for n in range(5, 9):
+                assert shifts_string[0 + 7 * n] == "-"
 
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
-                # balance_offset=10,
-                workers_required=1,
-            ),
-            SingleShift(
-                sites=("group1",), name="w", length= 12.5, days=days[5:],
-                workers_required=2,
-            ),
-        )
+def test_nwd_partial_overlay():
+    Rota = setup_rota()
+    start_date = Rota.start_date
+    d = start_date + datetime.timedelta(weeks=5)
+    d2 = start_date + datetime.timedelta(weeks=6)
+    worker1 = Worker(
+        name="worker1", site="group1", grade=1,
+        nwds=[{"day": "Mon", "start_date": start_date, "end_date": d2}],
+    )
+    worker2 = Worker(
+        name="worker2", site="group1", grade=1,
+        nwds=[{"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date}],
+    )
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    assert Rota.results.solver.status in ("warning", "error")
 
-        self.Rota.build_and_solve(options={"ratio": 0.10})
+def test_nwd_force_as_block():
+    Rota = setup_rota()
+    start_date = Rota.start_date
+    d = start_date + datetime.timedelta(weeks=4)
+    worker1 = Worker(
+        name="worker1", site="group1", grade=1,
+        nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
+    )
+    worker2 = Worker(name="worker2", site="group1", grade=1)
+    worker3 = Worker(name="worker3", site="group1", grade=1)
+    worker4 = Worker(name="worker4", site="group1", grade=1)
+    Rota.add_workers((worker1, worker2, worker3, worker4))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="d", length=12.5, days=days[:5],
+            balance_offset=10, workers_required=2,
+            force_as_block_unless_nwd=True, assign_as_block=True,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd", 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])
+        for week in weeks_from_list(shifts_string):
+            assert week in ("-------", "ddddd--")
+        assert shifts_string.count("d") == 25
 
-        self.Rota.export_rota_to_html("nwd1")
+def test_nwd_force_as_block_force_split():
+    Rota = setup_rota()
+    start_date = Rota.start_date
+    d = start_date + datetime.timedelta(weeks=5)
+    worker1 = Worker(
+        name="worker1", site="group1", grade=1,
+        nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
+    )
+    worker2 = Worker(name="worker2", site="group1", grade=1)
+    worker3 = Worker(name="worker3", site="group1", grade=1)
+    worker4 = Worker(name="worker4", site="group1", grade=1)
+    Rota.add_workers((worker1, worker2, worker3, worker4))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="d", length=12.5, days=days[:5],
+            workers_required=2, force_as_block_unless_nwd=True, assign_as_block=True,
+        ),
+        SingleShift(
+            sites=("group1",), name="w", length=12.5, days=days[5:],
+            workers_required=4, force_as_block_unless_nwd=True, assign_as_block=True,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd_block_force_split", folder="tests")
 
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
-
-        for worker in self.Rota.workers:
-            shifts = self.Rota.get_worker_shift_list(worker)
-
-            # Convert shift to a string representation
-            shifts_string = "".join([i if i != "" else "-" for i in shifts])
+    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])
+        for week in weeks_from_list(shifts_string[7 * 5 :]):
+            assert week in ("-----ww", "dddddww")
+        if worker.name == "worker1":
+            assert shifts_string[: 7 * 5].count("ww-") == 4
+        assert shifts_string.count("d") == 25
+        assert shifts_string.count("w") == 20
 
+def test_nwd_testing():
+    Rota = setup_rota()
+    start_date = Rota.start_date
+    d = start_date + datetime.timedelta(weeks=5)
+    worker1 = Worker(
+        name="worker1", site="group1", grade=1,
+        nwds=[
+            {"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date},
+            {"day": "Tue", "start_date": d, "end_date": Rota.rota_end_date},
+            {"day": "Wed", "start_date": d, "end_date": Rota.rota_end_date},
+            {"day": "Thu", "start_date": d, "end_date": Rota.rota_end_date},
+            {"day": "Fri", "start_date": d, "end_date": Rota.rota_end_date},
+        ],
+    )
+    worker2 = Worker(name="worker2", site="group1", grade=1)
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="d", length=12.5, days=days[:5],
+            balance_offset=10, workers_required=1,
+            force_as_block_unless_nwd=True, assign_as_block=True,
+        ),
+        SingleShift(
+            sites=("group1",), name="w", length=12.5, days=days[5:],
+            balance_offset=10, workers_required=2,
+            force_as_block_unless_nwd=True, assign_as_block=True,
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd", 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])
+        for week in weeks_from_list(shifts_string[7 * 5 :]):
             if worker.name == "worker1":
-                assert "wwd" not in shifts_string
-
-            if worker.name == "worker2":
-                assert "dww" not in shifts_string
-
-    def test_nwd_simple_fail(self):
-        # Set up rota
-        weeks_to_rota = 10
-        start_date = datetime.date(2022, 3, 7)
-
-        self.Rota = RotaBuilder(
-            start_date,
-            weeks_to_rota=weeks_to_rota,
-        )
-
-        # Add a few workers
-        worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
-        worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Mon"}])
-
-        self.Rota.add_workers((worker1, worker2))
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
-                balance_offset=10,
-                workers_required=1,
-            ),
-        )
-
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-
-        self.Rota.export_rota_to_html("nwd")
-
-        assert self.Rota.results.solver.status in ("warning", "error")
-
-    def test_nwd_partial_rota(self):
-        # Set up rota
-        weeks_to_rota = 10
-        start_date = datetime.date(2022, 3, 7)
-
-        self.Rota = RotaBuilder(
-            start_date,
-            weeks_to_rota=weeks_to_rota,
-        )
-
-        d = start_date + datetime.timedelta(weeks=5)
-        # Add a few workers
-        worker1 = Worker(
-            name="worker1",
-            site="group1",
-            grade=1,
-            nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
-        )
-        worker2 = Worker(
-            name="worker2",
-            site="group1",
-            grade=1,
-            nwds=[{"day": "Mon", "start_date": d, "end_date": self.Rota.rota_end_date}],
-        )
-
-        self.Rota.add_workers((worker1, worker2))
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
-                balance_offset=10,
-                workers_required=1,
-            ),
-        )
-
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-
-        self.Rota.export_rota_to_html("nwd")
-
-        assert self.Rota.results.solver.status == "ok"
-
-        for worker in self.Rota.workers:
-            shifts = self.Rota.get_worker_shift_list(worker)
-
-            # Convert shift to a string representation
-            shifts_string = "".join([i if i != "" else "-" for i in shifts])
-
+                assert week == "-----ww"
+            else:
+                assert week == "dddddww"
+        for week in weeks_from_list(shifts_string[: 7 * 5]):
             if worker.name == "worker1":
-                for n in range(0,4):
-                    assert shifts_string[0 + 7*n] == "-"
-
-            if worker.name == "worker2":
-                for n in range(5,9):
-                    assert shifts_string[0 + 7*n] == "-"
-
-
-    def test_nwd_partial_overlay(self):
-        # Set up rota
-        weeks_to_rota = 10
-        start_date = datetime.date(2022, 3, 7)
-
-        self.Rota = RotaBuilder(
-            start_date,
-            weeks_to_rota=weeks_to_rota,
-        )
-
-        d = start_date + datetime.timedelta(weeks=5)
-        d2 = start_date + datetime.timedelta(weeks=6)
-        # Add a few workers
-        worker1 = Worker(
-            name="worker1",
-            site="group1",
-            grade=1,
-            nwds=[{"day":"Mon", "start_date":start_date, "end_date":d2}],
-        )
-        worker2 = Worker(
-            name="worker2",
-            site="group1",
-            grade=1,
-            nwds=[{"day":"Mon", "start_date":d, "end_date":self.Rota.rota_end_date}],
-        )
-
-        self.Rota.add_workers((worker1, worker2))
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
-                balance_offset=10,
-                workers_required=1,
-            ),
-        )
-
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-
-        assert self.Rota.results.solver.status in ("warning", "error")
-
-    def test_nwd_force_as_block(self):
-        # Set up rota
-        weeks_to_rota = 10
-        start_date = datetime.date(2022, 3, 7)
-
-        self.Rota = RotaBuilder(
-            start_date,
-            weeks_to_rota=weeks_to_rota,
-        )
-
-        d = start_date + datetime.timedelta(weeks=4)
-        worker1 = Worker(
-            name="worker1",
-            site="group1",
-            grade=1,
-            nwds=[{"day": "Mon", "start_date":start_date, "end_date":d}],
-        )
-        worker2 = Worker(name="worker2", site="group1", grade=1,
-        )
-        worker3 = Worker(name="worker3", site="group1", grade=1,
-        )
-        worker4 = Worker(name="worker4", site="group1", grade=1,
-        )
-
-        self.Rota.add_workers((worker1, worker2, worker3, worker4))
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
-                balance_offset=10,
-                workers_required=2,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
-            ),
-        )
-
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.export_rota_to_html("nwd")
-
-        assert self.Rota.results.solver.status == "ok"
-
-        for worker in self.Rota.workers:
-            shifts = self.Rota.get_worker_shift_list(worker)
-
-            # Convert shift to a string representation
-            shifts_string = "".join([i if i != "" else "-" for i in shifts])
-
-            for week in weeks_from_list(shifts_string):
-                assert week in ("-------", "ddddd--")
-
-            assert shifts_string.count("d") == 25
-
-    def test_nwd_force_as_block_force_split(self):
-        # Set up rota
-        weeks_to_rota = 10
-        start_date = datetime.date(2022, 3, 7)
-
-        self.Rota = RotaBuilder(
-            start_date,
-            weeks_to_rota=weeks_to_rota,
-        )
-
-        d = start_date + datetime.timedelta(weeks=5)
-        worker1 = Worker(
-            name="worker1",
-            site="group1",
-            grade=1,
-            nwds=[{"day": "Mon", "start_date":start_date, "end_date":d}],
-        )
-        worker2 = Worker(name="worker2", site="group1", grade=1,
-        )
-        worker3 = Worker(name="worker3", site="group1", grade=1,
-        )
-        worker4 = Worker(name="worker4", site="group1", grade=1,
-        )
-
-        self.Rota.add_workers((worker1, worker2, worker3, worker4))
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
-                #balance_offset=50,
-                workers_required=2,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
-            ),
-            SingleShift(
-                sites=("group1",), name="w", length= 12.5, days=days[5:],
-                #balance_offset=50,
-                workers_required=4,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
-            ),
-        )
-
-        self.Rota.build_and_solve()
-        self.Rota.export_rota_to_html("nwd_block_force_split")
-
-        assert self.Rota.results.solver.status == "ok"
-
-        for worker in self.Rota.workers:
-            shifts = self.Rota.get_worker_shift_list(worker)
-
-            # Convert shift to a string representation
-            shifts_string = "".join([i if i != "" else "-" for i in shifts])
-
-            for week in weeks_from_list(shifts_string[7 * 5 :]):
-                assert week in ("-----ww", "dddddww")
-
-            if worker.name == "worker1":
-                assert shifts_string[: 7 * 5].count("ww-") == 4
-
-            assert shifts_string.count("d") == 25
-            assert shifts_string.count("w") == 20
-
-    def test_nwd_testing(self):
-        # Set up rota
-        weeks_to_rota = 10
-        start_date = datetime.date(2022, 3, 7)
-
-        self.Rota = RotaBuilder(
-            start_date,
-            weeks_to_rota=weeks_to_rota,
-        )
-
-        d = start_date + datetime.timedelta(weeks=5)
-        worker1 = Worker(name="worker1", site="group1", grade=1,
-            nwds=[
-                {"day": "Mon", "start_date":d, "end_date":self.Rota.rota_end_date},
-                {"day": "Tue", "start_date":d, "end_date":self.Rota.rota_end_date},
-                {"day": "Wed", "start_date":d, "end_date":self.Rota.rota_end_date},
-                {"day": "Thu", "start_date":d, "end_date":self.Rota.rota_end_date},
-                {"day": "Fri", "start_date":d, "end_date":self.Rota.rota_end_date},
-            ],
-        )
-        worker2 = Worker(name="worker2", site="group1", grade=1,
-        )
-
-        self.Rota.add_workers((worker1, worker2))
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
-                balance_offset=10,
-                workers_required=1,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
-            ),
-            SingleShift(
-                sites=("group1",), name="w", length= 12.5, days=days[5:],
-                balance_offset=10,
-                workers_required=2,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
-            ),
-        )
-
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.export_rota_to_html("nwd")
-
-        assert self.Rota.results.solver.status == "ok"
-
-        for worker in self.Rota.workers:
-            shifts = self.Rota.get_worker_shift_list(worker)
-
-            # Convert shift to a string representation
-            shifts_string = "".join([i if i != "" else "-" for i in shifts])
-
-            for week in weeks_from_list(shifts_string[7 * 5 :]):
-                if worker.name == "worker1":
-                    assert week == "-----ww"
-                else:
-                    assert week == "dddddww"
-
-            for week in weeks_from_list(shifts_string[: 7 * 5]):
-                if worker.name == "worker1":
-                    assert week == "dddddww"
-                else:
-                    assert week == "-----ww"
+                assert week == "dddddww"
+            else:
+                assert week == "-----ww"

test/test_rota_clear.py

@@ -0,0 +1,454 @@
+import pytest
+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
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+    )
+    Rota.constraint_options["max_shifts_per_month"] = 20
+
+    workers = [
+        Worker(name="worker1", site="group1", grade=1),
+        Worker(name="worker2", site="group1", grade=1),
+        Worker(name="worker3", site="group1", grade=1),
+        Worker(name="worker4", site="group1", grade=1),
+        Worker(name="worker5", site="group2", grade=1, fte=60),
+        Worker(name="worker6", site="group2", grade=1, fte=40),
+        Worker(name="worker7", site="group2", grade=1),
+        Worker(name="worker8", site="group2", grade=1),
+    ]
+    Rota.add_workers(workers)
+    return Rota
+
+
+def test_pre(demo_rota_clear):
+    Rota = demo_rota_clear
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="a",
+            length=12.5,
+            days=days[:5],
+            workers_required=2,
+            constraint=[{"name": "pre", "options": 1}],
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="b",
+            length=12.5,
+            days=days[:5],
+            workers_required=2,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="c",
+            length=12.5,
+            days=days[3],
+            workers_required=1,
+            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 1}],
+        ),
+    )
+
+    Rota.build_shifts()
+    Rota.build_workers()
+    Rota.build_model()
+
+    solver_options = {"seconds": 1000, "threads": 10}
+    Rota.solve_model(options=solver_options)
+    Rota.export_rota_to_html("test4")
+
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+
+    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 "ba" not in shifts_string
+        assert "ca" not in shifts_string
+        assert "ac" not in shifts_string
+        assert "ca" not in shifts_string
+        assert "bc" not in shifts_string
+        assert "cb" not in shifts_string
+        # Uncomment if you want to check for these as well:
+        # 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
+
+
+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.01})
+    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.01})
+    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

test/test_rota_nights.py

@@ -0,0 +1,99 @@
+import datetime
+import pytest
+from pytest import approx
+from rota.shifts import RotaBuilder, SingleShift, days
+from rota.workers import Worker
+
+@pytest.fixture
+def demo_rota_nights():
+    weeks_to_rota = 12
+    start_date = datetime.date(2022, 3, 14)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+    )
+
+    Rota.constraint_options["max_shifts_per_week"] = 5
+    Rota.constraint_options["max_shifts_per_month"] = 31
+    Rota.constraint_options["limit_to_1_st2_on_nights"] = False
+    Rota.constraint_options["constrain_time_off_after_nights"] = False
+    Rota.constraint_options["balance_nights_across_site"] = False
+    Rota.constraint_options["balance_shifts"] = True
+    Rota.constraint_options["balance_shifts_over_workers"] = True
+    Rota.constraint_options["balance_nights"] = False
+    Rota.constraint_options["minimise_shift_diffs"] = False
+    Rota.constraint_options["balance_blocks"] = False
+    Rota.constraint_options["max_night_frequency"] = 0
+
+    # Add workers
+    workers = [
+        Worker(name="worker1", site="group1", grade=1),
+        Worker(name="worker2", site="group1", grade=1),
+        Worker(name="worker3", site="group1", grade=1),
+        Worker(name="worker4", site="group1", grade=1),
+        Worker(name="worker5", site="group2", grade=1, fte=60),
+        Worker(name="worker6", site="group2", grade=1, fte=40),
+        Worker(name="worker7", site="group2", grade=1),
+        Worker(name="worker8", site="group2", grade=1),
+    ]
+    Rota.add_workers(workers)
+
+    # Add shifts
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1", "group2"),
+            name="night_weekday",
+            length=12.5,
+            days=days[:4],
+            workers_required=1,
+            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 2}],
+            force_as_block=True
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="night_weekend",
+            length=12.5,
+            days=days[4:],
+            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 2}],
+            force_as_block=True,
+        ),
+        SingleShift(
+            sites=("group1", "group2"),
+            name="twilight",
+            length=12.5,
+            days=days[:5],
+            workers_required=2,
+        ),
+    )
+
+    Rota.build_shifts()
+    Rota.build_workers()
+    Rota.build_model()
+
+    solver_options = {"ratio": 0.01, "seconds": 1000, "threads": 10}
+    Rota.solve_model(options=solver_options)
+    Rota.export_rota_to_html("test3")
+
+    return Rota
+
+def test_night_assignment(demo_rota_nights):
+    Rota = demo_rota_nights
+    assert Rota.results.solver.status == "ok"
+
+    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)
+
+        if worker.fte == 40:
+            assert worker_shifts["night_weekday"] in (approx(0), approx(4))
+            assert worker_shifts["night_weekend"] in (approx(0), approx(3))
+        elif worker.fte == 60:
+            assert worker_shifts["night_weekday"] in (approx(0), approx(4), approx(8))
+            assert worker_shifts["night_weekend"] in (approx(0), approx(3), approx(6))
+        else:
+            assert worker_shifts["night_weekday"] in (approx(4), approx(8))
+            assert worker_shifts["night_weekend"] in (approx(3), approx(6))

test/test_shift_constraints.py

@@ -0,0 +1,71 @@
+import pytest
+import datetime
+from copy import deepcopy
+from rota.shifts import RotaBuilder, SingleShift, days
+from rota.workers import Worker
+
+@pytest.fixture
+def rota_2workers():
+    weeks_to_rota = 10
+    start_date = datetime.date(2022, 3, 7)
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+    )
+    worker1 = Worker(name="worker1", site="group1", grade=1)
+    worker2 = Worker(name="worker2", site="group1", grade=1)
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=days,
+            balance_offset=40,
+            workers_required=1,
+        ),
+    )
+    return Rota
+
+@pytest.fixture
+def rota_3workers(rota_2workers):
+    Rota3 = deepcopy(rota_2workers)
+    worker3 = Worker(name="worker3", site="group1", grade=1)
+    Rota3.add_worker(worker3)
+    return Rota3
+
+def test_max_shifts(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_month"] = 14
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status == "ok"
+    assert rota_2workers.results.solver.termination_condition == "optimal"
+
+def test_max_shifts_fail(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_month"] = 13
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status in ("warning", "error")
+    assert rota_2workers.results.solver.termination_condition == "infeasible"
+
+def test_max_shifts_extra_worker(rota_3workers):
+    rota_3workers.constraint_options["max_shifts_per_month"] = 12
+    rota_3workers.build_and_solve()
+    assert rota_3workers.results.solver.status == "ok"
+    assert rota_3workers.results.solver.termination_condition == "optimal"
+
+def test_max_shifts_per_week_fail(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_week"] = 3
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status in ("warning", "error")
+    assert rota_2workers.results.solver.termination_condition == "infeasible"
+
+def test_max_shifts_per_week_pass(rota_2workers):
+    rota_2workers.constraint_options["max_shifts_per_month"] = 14
+    rota_2workers.constraint_options["max_shifts_per_week"] = 4
+    rota_2workers.build_and_solve()
+    assert rota_2workers.results.solver.status == "ok"
+
+def test_max_shifts_per_week_extra_worker_pass(rota_3workers):
+    rota_3workers.constraint_options["max_shifts_per_week"] = 3
+    rota_3workers.build_and_solve()
+    assert rota_3workers.results.solver.status == "ok"

test/test_shifts.py

@@ -708,4 +708,25 @@ class TestShiftWorkerRequirements:
         total_shifts = Rota.get_workers_total_shifts()
 
         for worker in Rota.get_workers():
-            assert total_shifts[worker.name] == 30
+            assert total_shifts[worker.name] == 30
+
+class TestDisplayChar:
+    def test_display_char_in_output(monkeypatch):
+        Rota = generate_basic_rota(workers=1)
+        shift = SingleShift(
+            sites=("group1",),
+            name="testshift",
+            length=8,
+            days=days[:1],
+            display_char="Z"
+        )
+        Rota.add_shifts(shift)
+
+        assert shift.display_char == "Z"
+
+        Rota.build_and_solve(options={"ratio": 0.1})
+
+        assert Rota.results.solver.status == "ok"
+
+        for worker in Rota.get_workers():
+            assert Rota.get_worker_shift_list_string(worker).startswith("Z------" * 5)

test/test_weekend_padding.py

@@ -46,7 +46,7 @@ class TestDemoRota:
 
         self.Rota.build_and_solve(options={"ratio": 0.00})
 
-        self.Rota.export_rota_to_html("balance_weekend")
+        self.Rota.export_rota_to_html("balance_weekend", folder="tests")
 
         assert self.Rota.results.solver.status == "ok"
         assert self.Rota.results.solver.termination_condition == "optimal"