feat: add exact shifts functionality to Worker class and update related logic in RotaBuilder

2026-07-08 22:33:33 +01:00
parent 02598cb665
commit 0f31dfe1a9
7 changed files with 398 additions and 102 deletions
@@ -1624,4 +1624,212 @@ def test_worker_hard_day_exclusion2():
        for week in range(16):
            sat_idx = week * 7 + days.index("Sat")
            sun_idx = week * 7 + days.index("Sun")
-            assert not (shift_list[sat_idx] == "a" and shift_list[sun_idx] == "a"), f"Worker {worker.name} has 'a' on both Mon and Tue in week {week}, which should not happen due to hard day exclusion"
+            assert not (shift_list[sat_idx] == "a" and shift_list[sun_idx] == "a"), f"Worker {worker.name} has 'a' on both Mon and Tue in week {week}, which should not happen due to hard day exclusion"
+
+
+def test_exact_shifts_distribution():
+    weeks_to_rota = 10
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    workers = [
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"a": 4}),
+        Worker(name="worker2", site="group1", grade=1, fte=100),
+        Worker(name="worker3", site="group1", grade=1, fte=50),
+    ]
+    Rota.add_workers(workers)
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            hard_constrain_shift=True,
+            balance_offset=0,
+        ),
+    )
+
+    Rota.build_and_solve()
+
+    assert Rota.workers_name_map["worker1"].shift_target_number["a"] == 4
+    assert Rota.workers_name_map["worker2"].shift_target_number["a"] == 4
+    assert Rota.workers_name_map["worker3"].shift_target_number["a"] == 2
+
+    # Verify actual assigned counts are exactly correct
+    worker_shift_counts = {w.name: 0 for w in Rota.get_workers()}
+    for week, day, shiftname in Rota.get_all_shiftname_combinations():
+        for worker in Rota.get_workers():
+            val = Rota.model.works[worker.id, week, day, shiftname].value
+            if val is not None and val > 0.5:
+                if shiftname == "a":
+                    worker_shift_counts[worker.name] += 1
+
+    assert worker_shift_counts["worker1"] == 4
+    assert worker_shift_counts["worker2"] == 4
+    assert worker_shift_counts["worker3"] == 2
+
+
+def test_exact_shifts_distribution_unbalanced():
+    weeks_to_rota = 16
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    workers = [
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"a": 4}),
+        Worker(name="worker2", site="group1", grade=1, fte=100),
+        Worker(name="worker3", site="group1", grade=1, fte=50),
+    ]
+    Rota.add_workers(workers)
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            hard_constrain_shift=True,
+            balance_offset=0,
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=12.5,
+            days=("Tue",),
+            workers_required=1,
+            hard_constrain_shift=True,
+            balance_offset=0.8,
+        ),
+    )
+
+    Rota.build_and_solve()
+    Rota.export_rota_to_html("test_exact_shifts_distribution_unbalanced", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+
+    assert Rota.workers_name_map["worker1"].shift_target_number["a"] == 4
+    assert Rota.workers_name_map["worker2"].shift_target_number["a"] == 8
+    assert Rota.workers_name_map["worker3"].shift_target_number["a"] == 4
+
+    assert Rota.workers_name_map["worker1"].shift_target_number["b"] == 6.4
+    assert Rota.workers_name_map["worker2"].shift_target_number["b"] == 6.4
+    assert Rota.workers_name_map["worker3"].shift_target_number["b"] == 3.2
+
+    # Verify actual assigned counts are exactly correct
+    worker_shift_counts_a = {w.name: 0 for w in Rota.get_workers()}
+    worker_shift_counts_b = {w.name: 0 for w in Rota.get_workers()}
+    for week, day, shiftname in Rota.get_all_shiftname_combinations():
+        for worker in Rota.get_workers():
+            val = Rota.model.works[worker.id, week, day, shiftname].value
+            if val is not None and val > 0.5:
+                if shiftname == "a":
+                    worker_shift_counts_a[worker.name] += 1
+                elif shiftname == "b":
+                    worker_shift_counts_b[worker.name] += 1
+
+    assert worker_shift_counts_a["worker1"] == 4
+    assert worker_shift_counts_a["worker2"] == 8
+    assert worker_shift_counts_a["worker3"] == 4
+
+    assert worker_shift_counts_b["worker1"] in (6, 7)  # Allowing for rounding differences
+    assert worker_shift_counts_b["worker2"] in (6, 7)  # Allowing for rounding differences
+    assert worker_shift_counts_b["worker3"] in (3, 4)  # Allowing for rounding differences
+
+
+def test_exact_shifts_invalid_site_warning():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group2", grade=1, fte=100, exact_shifts={"a": 4}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+        SingleShift(
+            sites=("group2",),
+            name="b",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+    )
+    Rota.build_workers()
+    
+    warnings = Rota.get_warnings("Invalid exact shift site")
+    assert len(warnings) > 0
+    assert "group2 is not in the shift sites" in warnings[0][1]
+
+
+def test_exact_shifts_invalid_shift_warning():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"nonexistent": 4}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+    )
+    Rota.build_workers()
+    
+    warnings = Rota.get_warnings("Invalid exact shift")
+    assert len(warnings) > 0
+    assert "non-existent shift nonexistent" in warnings[0][1]
+
+
+def test_exact_shifts_exceed_total_warning():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"a": 4}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+    )
+    Rota.build_workers()
+    Rota.build_model()
+    
+    warnings = Rota.get_warnings("Exact shifts exceed total shifts")
+    assert len(warnings) > 0
+    assert "total exact shifts (4) exceeds total required shifts" in warnings[0][1]