major test refactoring

2025-06-05 23:40:28 +01:00
parent b9ca76aea1
commit a647416720
13 changed files with 1690 additions and 2282 deletions
@@ -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"