major test refactoring

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

rota/shifts.py

@@ -392,6 +392,7 @@ class RotaBuilder(object):
         debug_if_fail: bool = False,
     ):
         console.print("Setting up solver")
+        console.print(f"{options=}")
 
         if solver == "scip":
             if "seconds" in options:
@@ -516,6 +517,9 @@ class RotaBuilder(object):
         self.build_workers()
         self.build_model()
 
+        console.print("Building model")
+        console.print(f"{options=}")
+
         if solve:
             self.solve_model(
                 options=options, debug_if_fail=debug_if_fail, solver=solver
@@ -2203,7 +2207,6 @@ class RotaBuilder(object):
                             self.get_week_start_date(week)
                         ):
                             # Force shifts to be assigned in blocks
-                            try:
                                 self.model.constraints.add(
                                     self.model.shift_week_worker_assigned[
                                         shift.name, week, worker.id
@@ -2214,8 +2217,6 @@ class RotaBuilder(object):
                                         for day in shift.days
                                     )
                                 )
-                            except KeyError:
-                                pass # blocks with workerd requirements
 
                 #
 
@@ -3213,7 +3214,7 @@ class RotaBuilder(object):
             shifts = []
 
             for shift in self.shifts:
-                if shift.start_date < date <= shift.end_date:
+                if shift.start_date <= date < shift.end_date:
                     if shift.get_worker_requirement_by_date(date):
                         shifts.append(shift)
 

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,39 +1,27 @@
 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():
-class TestBalancing:
-    def test_basic_balancing(self):
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -54,21 +42,16 @@ class TestBalancing:
             force_as_block=False,
         ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.001})
     Rota.export_rota_to_html("basic_balancing")
-
     shift_summary = Rota.get_shift_summary_dict()
     for worker_name in shift_summary:
         worker_shifts = shift_summary[worker_name]
-
-            #worker = Rota.get_worker_by_name(worker_name)
         assert worker_shifts["a"] in (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):
+def test_weighted_shift_balancing():
     Rota = generate_basic_rota(20)
-
     Rota.add_shifts(
         SingleShift(
             sites=["group1", "group2"],
@@ -97,32 +80,23 @@ class TestBalancing:
             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):
+def test_weighted_shift_balancing2():
     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"),
@@ -142,32 +116,23 @@ class TestBalancing:
             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):
+def test_weighted_shift_balancing3():
     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"),
@@ -188,26 +153,18 @@ class TestBalancing:
             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):
+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"),
@@ -218,7 +175,6 @@ class TestBalancing:
             balance_weighting=8,
             workers_required=2,
             force_as_block=False,
-                # constraint=["preclear2", "postclear2"],
         ),
         SingleShift(
             sites=("group2", "group3"),
@@ -226,29 +182,22 @@ class TestBalancing:
             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):
+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"),
@@ -256,10 +205,8 @@ class TestBalancing:
             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"),
@@ -272,13 +219,9 @@ class TestBalancing:
             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)

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: 2}}],
+        ),
+    )
+    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 = 2 if grade == 3 else 2
+        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_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,44 +1,33 @@
 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():
-class TestNightShifts:
-    def test_nights(self):
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -55,19 +44,15 @@ class TestNightShifts:
             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")
+    Rota.export_rota_to_html("night", folder="tests")
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
 
-    def test_nights_pre3(self):
+def test_nights_pre3():
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -84,18 +69,14 @@ class TestNightShifts:
             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")
+    Rota.export_rota_to_html("night_pre3", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_pre3_2(self):
+def test_nights_pre3_2():
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -112,18 +93,14 @@ class TestNightShifts:
             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")
+    Rota.export_rota_to_html("night_pre3", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_fail(self):
+def test_nights_fail():
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -140,19 +117,14 @@ class TestNightShifts:
             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):
+def test_nights2():
     Rota = generate_basic_rota()
-
     Rota.constraint_options["max_night_frequency"] = 0
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -169,27 +141,16 @@ class TestNightShifts:
             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")
+    Rota.export_rota_to_html("night", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_pre_wrap_around(self):
+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"),
@@ -204,20 +165,7 @@ class TestNightShifts:
             ],
             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.add_workers(
         [
             Worker(
@@ -229,25 +177,17 @@ class TestNightShifts:
                         "date": Rota.rota_end_date - datetime.timedelta(days=1),
                         "reason": "no wanna",
                     },
-                        # {
-                        #    "date": Rota.rota_end_date - datetime.timedelta(days=5),
-                        #    "reason": "no wanna",
-                        # },
                 ],
             ),
         ]
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
+    Rota.export_rota_to_html("night pre wrap around", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_max_frequency(self):
+def test_nights_max_frequency():
     Rota = generate_basic_rota()
-
     Rota.constraint_options["max_night_frequency"] = 1
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -258,33 +198,14 @@ class TestNightShifts:
             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")
+    Rota.export_rota_to_html("night max frequency", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_max_frequency_fail(self):
+def test_nights_max_frequency_fail():
     Rota = generate_basic_rota()
-
     Rota.constraint_options["max_night_frequency"] = 2
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -295,33 +216,13 @@ class TestNightShifts:
             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 in ("warning", "error")
 
-    def test_nights_max_frequency3(self):
+def test_nights_max_frequency3():
     Rota = generate_basic_rota()
-
     Rota.constraint_options["max_night_frequency"] = 2
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -338,27 +239,17 @@ class TestNightShifts:
             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.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
+    Rota.export_rota_to_html("night max frequency 3", folder="tests")
-
     assert Rota.results.solver.status in ("warning", "error")
 
-    def test_nights_max_frequency4(self):
+def test_nights_max_frequency4():
     Rota = generate_basic_rota()
-
     Rota.constraint_options["max_night_frequency"] = 2
-
     Rota.add_shifts(
         SingleShift(
             sites=("group1", "group2"),
@@ -375,43 +266,24 @@ class TestNightShifts:
             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(
+            Worker(name="worker3", site="group1", grade=1),
-                    name="worker3",
+            Worker(name="worker4", site="group1", grade=1),
-                    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")
+    Rota.export_rota_to_html("night max frequency 4", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_max_frequency_exclusions(self):
+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"),
@@ -423,18 +295,14 @@ class TestNightShifts:
             workers_required=2,
         ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
+    Rota.export_rota_to_html("night max frequency exclusions", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_max_frequency_exclusions2(self):
+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"),
@@ -446,7 +314,6 @@ class TestNightShifts:
             workers_required=2,
         ),
     )
-
     Rota.add_workers(
         [
             Worker(
@@ -463,19 +330,14 @@ class TestNightShifts:
             ),
         ]
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
+    Rota.export_rota_to_html("night max frequency exclusions 2", folder="tests")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_nights_max_frequency_exclusions3(self):
+def test_nights_max_frequency_exclusions3():
-        # 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"),
@@ -488,7 +350,6 @@ class TestNightShifts:
             workers_required=2,
         ),
     )
-
     Rota.add_workers(
         [
             Worker(
@@ -505,8 +366,6 @@ class TestNightShifts:
             ),
         ]
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-        Rota.export_rota_to_html("night")
+    Rota.export_rota_to_html("night max frequency exclusions 3", folder="tests")
-
     assert Rota.results.solver.status == "ok"

test/test_nwds.py

@@ -1,372 +1,214 @@
 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]
 
-
+def setup_rota(weeks_to_rota=10, start_date=None):
-class TestDemoRota:
+    if start_date is None:
-    def test_nwd(self):
-        # Set up rota
-        weeks_to_rota = 10
         start_date = datetime.date(2022, 3, 7)
+    return RotaBuilder(start_date, weeks_to_rota=weeks_to_rota)
 
-        self.Rota = RotaBuilder(
+def test_nwd():
-            start_date,
+    Rota = setup_rota()
-            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": "Fri"}])
-
+    Rota.add_workers((worker1, worker2))
-        self.Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
         SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
+            sites=("group1",), name="d", length=12.5, days=days[:5], workers_required=1,
-                # balance_offset=10,
-                workers_required=1,
         ),
         SingleShift(
-                sites=("group1",), name="w", length= 12.5, days=days[5:],
+            sites=("group1",), name="w", length=12.5, days=days[5:], workers_required=2,
-                workers_required=2,
         ),
     )
-
+    Rota.build_and_solve(options={"ratio": 0.10})
-        self.Rota.build_and_solve(options={"ratio": 0.10})
+    Rota.export_rota_to_html("nwd1", folder="tests")
-
+    assert Rota.results.solver.status == "ok"
-        self.Rota.export_rota_to_html("nwd1")
+    assert Rota.results.solver.termination_condition == "optimal"
-
+    for worker in Rota.workers:
-        assert self.Rota.results.solver.status == "ok"
+        shifts = Rota.get_worker_shift_list(worker)
-        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])
-
         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):
+def test_nwd_simple_fail():
-        # Set up rota
+    Rota = setup_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"}])
-
+    Rota.add_workers((worker1, worker2))
-        self.Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
         SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
+            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
-                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.build_and_solve(options={"ratio": 0.00})
+def test_nwd_partial_rota():
-
+    Rota = setup_rota()
-        self.Rota.export_rota_to_html("nwd")
+    start_date = Rota.start_date
-
-        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",
+        name="worker1", site="group1", grade=1,
-            site="group1",
-            grade=1,
         nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
     )
     worker2 = Worker(
-            name="worker2",
+        name="worker2", site="group1", grade=1,
-            site="group1",
+        nwds=[{"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date}],
-            grade=1,
-            nwds=[{"day": "Mon", "start_date": d, "end_date": self.Rota.rota_end_date}],
     )
-
+    Rota.add_workers((worker1, worker2))
-        self.Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
         SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
+            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
-                balance_offset=10,
-                workers_required=1,
         ),
     )
-
+    Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd", folder="tests")
-
+    assert Rota.results.solver.status == "ok"
-        self.Rota.export_rota_to_html("nwd")
+    for worker in Rota.workers:
-
+        shifts = Rota.get_worker_shift_list(worker)
-        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])
-
         if worker.name == "worker1":
-                for n in range(0,4):
+            for n in range(0, 4):
-                    assert shifts_string[0 + 7*n] == "-"
+                assert shifts_string[0 + 7 * n] == "-"
-
         if worker.name == "worker2":
-                for n in range(5,9):
+            for n in range(5, 9):
-                    assert shifts_string[0 + 7*n] == "-"
+                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,
-        )
 
+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)
-        # Add a few workers
     worker1 = Worker(
-            name="worker1",
+        name="worker1", site="group1", grade=1,
-            site="group1",
+        nwds=[{"day": "Mon", "start_date": start_date, "end_date": d2}],
-            grade=1,
-            nwds=[{"day":"Mon", "start_date":start_date, "end_date":d2}],
     )
     worker2 = Worker(
-            name="worker2",
+        name="worker2", site="group1", grade=1,
-            site="group1",
+        nwds=[{"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date}],
-            grade=1,
-            nwds=[{"day":"Mon", "start_date":d, "end_date":self.Rota.rota_end_date}],
     )
-
+    Rota.add_workers((worker1, worker2))
-        self.Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
-
-        # Add a weekday and weekend shift
-        self.Rota.add_shifts(
         SingleShift(
-                sites=("group1",), name="d", length= 12.5, days=days[:5],
+            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
-                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.00})
+def test_nwd_force_as_block():
-
+    Rota = setup_rota()
-        assert self.Rota.results.solver.status in ("warning", "error")
+    start_date = Rota.start_date
-
-    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",
+        name="worker1", site="group1", grade=1,
-            site="group1",
+        nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
-            grade=1,
-            nwds=[{"day": "Mon", "start_date":start_date, "end_date":d}],
     )
-        worker2 = Worker(name="worker2", site="group1", grade=1,
+    worker2 = Worker(name="worker2", site="group1", grade=1)
-        )
+    worker3 = Worker(name="worker3", 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))
-        worker4 = Worker(name="worker4", site="group1", grade=1,
+    Rota.add_shifts(
-        )
-
-        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],
+            sites=("group1",), name="d", length=12.5, days=days[:5],
-                balance_offset=10,
+            balance_offset=10, workers_required=2,
-                workers_required=2,
+            force_as_block_unless_nwd=True, assign_as_block=True,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
         ),
     )
-
+    Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd", folder="tests")
-        self.Rota.export_rota_to_html("nwd")
+    assert Rota.results.solver.status == "ok"
-
+    for worker in Rota.workers:
-        assert self.Rota.results.solver.status == "ok"
+        shifts = Rota.get_worker_shift_list(worker)
-
-        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):
+def test_nwd_force_as_block_force_split():
-        # Set up rota
+    Rota = setup_rota()
-        weeks_to_rota = 10
+    start_date = Rota.start_date
-        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",
+        name="worker1", site="group1", grade=1,
-            site="group1",
+        nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
-            grade=1,
-            nwds=[{"day": "Mon", "start_date":start_date, "end_date":d}],
     )
-        worker2 = Worker(name="worker2", site="group1", grade=1,
+    worker2 = Worker(name="worker2", site="group1", grade=1)
-        )
+    worker3 = Worker(name="worker3", 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))
-        worker4 = Worker(name="worker4", site="group1", grade=1,
+    Rota.add_shifts(
-        )
-
-        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],
+            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,
-                workers_required=2,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
         ),
         SingleShift(
-                sites=("group1",), name="w", length= 12.5, days=days[5:],
+            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,
-                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")
 
-        self.Rota.build_and_solve()
+    assert Rota.results.solver.status == "ok"
-        self.Rota.export_rota_to_html("nwd_block_force_split")
+    for worker in Rota.workers:
-
+        shifts = Rota.get_worker_shift_list(worker)
-        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):
+def test_nwd_testing():
-        # Set up rota
+    Rota = setup_rota()
-        weeks_to_rota = 10
+    start_date = Rota.start_date
-        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,
+    worker1 = Worker(
+        name="worker1", site="group1", grade=1,
         nwds=[
-                {"day": "Mon", "start_date":d, "end_date":self.Rota.rota_end_date},
+            {"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date},
-                {"day": "Tue", "start_date":d, "end_date":self.Rota.rota_end_date},
+            {"day": "Tue", "start_date": d, "end_date": Rota.rota_end_date},
-                {"day": "Wed", "start_date":d, "end_date":self.Rota.rota_end_date},
+            {"day": "Wed", "start_date": d, "end_date": Rota.rota_end_date},
-                {"day": "Thu", "start_date":d, "end_date":self.Rota.rota_end_date},
+            {"day": "Thu", "start_date": d, "end_date": Rota.rota_end_date},
-                {"day": "Fri", "start_date":d, "end_date":self.Rota.rota_end_date},
+            {"day": "Fri", "start_date": d, "end_date": Rota.rota_end_date},
         ],
     )
-        worker2 = Worker(name="worker2", site="group1", grade=1,
+    worker2 = Worker(name="worker2", site="group1", grade=1)
-        )
+    Rota.add_workers((worker1, worker2))
-
+    Rota.add_shifts(
-        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],
+            sites=("group1",), name="d", length=12.5, days=days[:5],
-                balance_offset=10,
+            balance_offset=10, workers_required=1,
-                workers_required=1,
+            force_as_block_unless_nwd=True, assign_as_block=True,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
         ),
         SingleShift(
-                sites=("group1",), name="w", length= 12.5, days=days[5:],
+            sites=("group1",), name="w", length=12.5, days=days[5:],
-                balance_offset=10,
+            balance_offset=10, workers_required=2,
-                workers_required=2,
+            force_as_block_unless_nwd=True, assign_as_block=True,
-                force_as_block_unless_nwd=True,
-                assign_as_block=True,
         ),
     )
-
+    Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.build_and_solve(options={"ratio": 0.00})
+    Rota.export_rota_to_html("nwd", folder="tests")
-        self.Rota.export_rota_to_html("nwd")
+    assert Rota.results.solver.status == "ok"
-
+    for worker in Rota.workers:
-        assert self.Rota.results.solver.status == "ok"
+        shifts = Rota.get_worker_shift_list(worker)
-
-        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"

test/test_rota_clear.py

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

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"