feat: add group management for workers and shifts, including validation for overlapping groups

test/test_workers.py

@@ -1,5 +1,6 @@
 import pytest
 from rota_generator.shifts import InvalidShift, MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
+from pydantic import ValidationError
 
 import datetime
 from rota_generator.workers import NonWorkingDays, NotAvailableToWork, Worker, generate_not_available_to_works
@@ -1868,4 +1869,119 @@ def test_html_export_worker_details_and_unavailable_reason():
     html = Rota.get_worker_timetable_html()
     
     assert 'data-exact-shifts=\'{"a": 1}\'' in html
-    assert 'title=\' (2022-03-07) / START DATE: 2022-03-14\'' in html
+    assert 'title=\' (2022-03-07) / START DATE: 2022-03-14\'' in html
+
+
+def test_worker_groups_selector():
+    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, groups={"groupA", "groupB"}),
+        Worker(name="worker2", site="group1", grade=1, fte=100, groups={"groupB", "groupC"}),
+    ])
+    
+    workers_A = Rota.get_workers_in_group("groupA")
+    workers_B = Rota.get_workers_in_group("groupB")
+    workers_C = Rota.get_workers_in_group("groupC")
+    workers_D = Rota.get_workers_in_group("groupD")
+    
+    assert [w.name for w in workers_A] == ["worker1"]
+    assert set(w.name for w in workers_B) == {"worker1", "worker2"}
+    assert [w.name for w in workers_C] == ["worker2"]
+    assert workers_D == []
+
+
+def test_shift_groups_overlapping_validation():
+    with pytest.raises(ValidationError) as exc_info:
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            include_groups={"pool1", "pool2"},
+            exclude_groups={"pool2", "pool3"},
+        )
+    assert "include_groups and exclude_groups cannot overlap" in str(exc_info.value)
+
+
+def test_shift_group_availability_inclusion():
+    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),
+        Worker(name="worker2", site="group2", grade=1, fte=100, groups={"senior"}),
+        Worker(name="worker3", site="group2", grade=1, fte=100, groups={"junior"}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            include_groups={"senior"},
+        ),
+    )
+    Rota.terminate_on_warning.remove("Worker/no valid shifts")
+    Rota.build_and_solve()
+    Rota.export_rota_to_html("test_shift_group_availability_inclusion", folder="tests")
+    assert Rota.results.solver.status == "ok"
+    
+    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:
+                worker_shift_counts[worker.name] += 1
+                
+    assert worker_shift_counts["worker1"] in (1, 2)
+    assert worker_shift_counts["worker2"] in (1, 2)
+    assert worker_shift_counts["worker3"] == 0
+
+
+def test_shift_group_availability_exclusion():
+    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, groups={"senior"}),
+        Worker(name="worker2", site="group1", grade=1, fte=100, groups={"junior"}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            exclude_groups={"junior"},
+        ),
+    )
+    Rota.terminate_on_warning.remove("Worker/no valid shifts")
+    Rota.build_and_solve()
+    assert Rota.results.solver.status == "ok"
+    
+    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:
+                worker_shift_counts[worker.name] += 1
+                
+    assert worker_shift_counts["worker1"] == 2
+    assert worker_shift_counts["worker2"] == 0