Add forced shift assignment functionality for workers and update tests

test/test_workers.py

@@ -1455,4 +1455,83 @@ def test_max_shifts_per_week_by_shift_name_invalid_shift():
         ),
     )
     with pytest.raises(InvalidShift):
-        Rota.build_and_solve(options={"ratio": 0.0})
+        Rota.build_and_solve(options={"ratio": 0.0})
+
+def test_force_assign_shift():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=2)
+
+    workers = Rota.get_workers()
+
+    workers[0].force_assign_shift(1, "Mon", "a")
+    workers[0].force_assign_shift(1, "Tue", "b")
+    workers[0].force_assign_shift(2, "Wed", "b")
+    workers[0].force_assign_shift(2, "Thu", "b")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign", folder="tests")
+
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        if worker == workers[0]:
+            # Worker 0 should have 'a' on Monday and 'b' on Tuesday
+            assert shift_list[0] == "a", "Worker 0 should have 'a' on 1st Monday"
+            assert shift_list[1] == "b", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[9] == "b", "Worker 0 should have 'b' on 2nd Wednesday"
+            assert shift_list[10] == "b", "Worker 0 should have 'b' on 2nd Thursday"
+
+def test_force_assign_shift2():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=1)
+
+    workers = Rota.get_workers()
+
+    workers[0].force_assign_shift(1, "Mon", "a")
+    workers[0].force_assign_shift(1, "Tue", "a")
+    workers[0].force_assign_shift(1, "Wed", "a")
+    workers[0].force_assign_shift(1, "Thu", "a")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign", folder="tests")
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        if worker == workers[0]:
+            # Worker 0 should have 'a' on Monday and 'b' on Tuesday
+            assert shift_list[0] == "a", "Worker 0 should have 'a' on 1st Monday"
+            assert shift_list[1] == "a", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[2] == "a", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[3] == "a", "Worker 0 should have 'b' on 1st Tuesday"