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,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))