proc-rota/test/test_workers.py

import pytest
from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days

import datetime
from rota.workers import NonWorkingDays, NotAvailableToWork, Worker, generate_not_available_to_works
from rota.workers import WorkRequests

def generate_basic_rota(
    weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)
):
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

    # Add a few workers
    if workers:
        Rota.add_workers(
            [
                Worker(name=f"worker{i:02}", site="group1", grade=1)
                for i in range(1, workers + 1)
            ]
        )

    return Rota

def test_no_workers():
    """Check error raised if attempt to solve with no workers"""
    Rota = generate_basic_rota()

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="d",
            length=12.5,
            days=days[:5],
            balance_offset=10,
            workers_required=1,
        ),
        SingleShift(
            sites=("group1",),
            name="w",
            length=12.5,
            days=days[5:],
            workers_required=1,
        ),
    )

    with pytest.raises(NoWorkers):
        Rota.build_and_solve(solve=False)

def test_workers_duplicate_name():
    Rota = generate_basic_rota()

    workers = []
    workers.append(Worker(name="worker1", site="group1", grade=1, fte=100))
    workers.append(Worker(name="worker1", site="group1", grade=1, fte=100))
    workers.append(Worker(name="worker3", site="group2", grade=1, fte=50))
    workers.append(Worker(name="worker4", site="group2", grade=1, fte=50))

    Rota.add_workers(workers)

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="d",
            length=12.5,
            days=days[:5],
            balance_offset=10,
            workers_required=1,
        ),
    )

    with pytest.raises(WarningTermination):
        Rota.build_and_solve(solve=False)

def test_worker_ftes():
    weeks_to_rota = 10
    start_date = datetime.date(2022, 3, 7)

    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

    worker1_end_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 2)
    worker2_start_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 2)
    worker3_start_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 4)
    worker3_end_date = start_date + datetime.timedelta(weeks=3 * weeks_to_rota / 4)

    workers = []
    workers.append(
        Worker(
            name="worker1",
            site="group1",
            grade=1,
            fte=100,
            end_date=worker1_end_date,
        )
    )
    workers.append(
        Worker(
            name="worker2",
            site="group1",
            grade=1,
            fte=100,
            start_date=worker2_start_date,
        )
    )
    workers.append(
        Worker(
            name="worker3",
            site="group2",
            grade=1,
            fte=100,
            start_date=worker3_start_date,
            end_date=worker3_end_date,
        )
    )
    workers.append(Worker(name="worker4", site="group2", grade=1, fte=50))
    workers.append(
        Worker(name="worker5", site="group2", grade=1, fte=50, nwd=["Mon", "Tue"])
    )
    workers.append(
        Worker(
            name="worker6",
            site="group3",
            grade=1,
            fte=100,
            start_date=start_date + datetime.timedelta(weeks=weeks_to_rota),
        )
    )
    workers.append(
        Worker(name="worker7", site="group3", grade=1, fte=50, end_date=start_date)
    )

    Rota.add_workers(workers)

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="d",
            length=12.5,
            days=days[:5],
            balance_offset=10,
            workers_required=1,
        ),
    )

    with pytest.raises(WarningTermination):
        Rota.build_and_solve()

    Rota.terminate_on_warning.remove("Worker/no valid shifts")
    Rota.build_and_solve()

    # All workers should have an adjusted fte of 50 %
    for worker in Rota.get_workers():
        # Except those in group3
        if worker.site == "group3":
            assert worker.fte_adj == 0
            assert Rota.get_worker_shift_list(worker).count("d") == 0
            continue

        assert worker.fte_adj == 50

        if worker.fte == 100:
            assert worker.proportion_rota_to_work == 0.5
            assert worker.days_to_work == 35

        assert Rota.get_worker_shift_list(worker).count("d") == 10

        if worker.name == "worker1":
            assert Rota.get_worker_shift_list_string(worker).endswith(
                "-" * (Rota.rota_days_length // 2)
            )
        if worker.name == "worker2":
            assert Rota.get_worker_shift_list_string(worker).startswith(
                "-" * (Rota.rota_days_length // 2)
            )
        if worker.name == "worker3":
            assert Rota.get_worker_shift_list_string(worker).startswith(
                "-" * (Rota.rota_days_length // 4)
            )
            assert Rota.get_worker_shift_list_string(worker).endswith(
                "-" * (Rota.rota_days_length // 4)
            )

    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

    Rota.export_rota_to_html("worker_ftes")

def test_worker_oops():
    weeks_to_rota = 10
    start_date = datetime.date(2022, 3, 7)

    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

    worker1_oop = [
        {
            "start_date": start_date,
            "end_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
            "reason": "First half OOP",
        }
    ]

    worker2_oop = [
        {
            "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
            "end_date": Rota.rota_end_date,
            "reason": "First half OOP",
        }
    ]

    worker3_start_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 4)
    worker3_oop = [
        {
            "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
            "end_date": start_date
            + datetime.timedelta(weeks=3 * weeks_to_rota / 4),
            "reason": "Quarter oop",
        }
    ]

    worker4_oop = [
        {
            "start_date": start_date,
            "end_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 4),
            "reason": "Quarter oop",
        },
        {
            "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
            "end_date": start_date
            + datetime.timedelta(weeks=weeks_to_rota / 2)
            + datetime.timedelta(weeks=weeks_to_rota / 4),
            "reason": "Quarter oop 2",
        },
    ]

    workers = []
    workers.append(
        Worker(name="worker1", site="group1", grade=1, fte=100, oop=worker1_oop)
    )
    workers.append(
        Worker(name="worker2", site="group1", grade=1, fte=100, oop=worker2_oop)
    )
    workers.append(
        Worker(
            name="worker3",
            site="group2",
            grade=1,
            fte=100,
            start_date=worker3_start_date,
            oop=worker3_oop,
        )
    )
    workers.append(Worker(name="worker4", site="group2", grade=1, oop=worker4_oop))
    workers.append(
        Worker(name="worker5", site="group2", grade=1, fte=50, nwd=["Mon", "Tue"])
    )

    worker6_oop = [
        {
            "start_date": start_date,
            "end_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 6),
            "reason": "oop",
        },
        {
            "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
            "end_date": start_date
            + datetime.timedelta(weeks=weeks_to_rota / 2)
            + datetime.timedelta(weeks=weeks_to_rota / 3),
            "reason": " oop 2",
        },
    ]

    worker7_oop = [
        {
            "start_date": start_date - datetime.timedelta(weeks=1),
            "end_date": Rota.rota_end_date,
            "reason": "oop",
        },
    ]

    workers.append(
        Worker(name="worker6", site="group2", grade=1, fte=100, oop=worker6_oop)
    )
    workers.append(
        Worker(name="worker7", site="group3", grade=1, fte=100, oop=worker7_oop)
    )

    Rota.add_workers(workers)

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="d",
            length=12.5,
            days=days[:5],
            balance_offset=10,
            workers_required=1,
        ),
    )

    Rota.terminate_on_warning.remove("Worker/no valid shifts")
    Rota.build_and_solve()

    # All workers should have an adjusted fte of 50 %
    for worker in Rota.get_workers():
        if worker.site == "group3":
            assert worker.fte_adj == 0
            continue

        assert worker.fte_adj == pytest.approx(50, 2)

        if worker.fte == 100:
            assert worker.proportion_rota_to_work == pytest.approx(0.5, 0.2)
            assert worker.days_to_work == pytest.approx(35, 1)

        assert Rota.get_worker_shift_list(worker).count("d") in (8, 9)

    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

    Rota.export_rota_to_html("worker_oops")

def test_worker_pairs():
    Rota = generate_basic_rota()

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=12.5,
            days=days[:5],
            balance_offset=20,
            workers_required=1,
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=12.5,
            days=days[:5],
            balance_offset=20,
            workers_required=1,
        ),
    )

    Rota.add_workers(
        [
            Worker(name="worker1", site="group1", grade=1, fte=100, pair="1"),
            Worker(name="worker2", site="group1", grade=1, fte=100, pair="1"),
        ]
    )

    Rota.constraint_options["hard_constrain_pair_separation"] = True

    Rota.build_and_solve()

    # With two paired workers we shouldn't be able to solve
    assert Rota.results.solver.status in ("warning", "error")

    worker3 = Worker(name="worker3", site="group1", grade=1, fte=100)
    Rota.add_worker(worker3)

    # An additional worker can take half the shifts
    Rota.build_and_solve()

    assert Rota.results.solver.status == "ok"

    assert Rota.get_workers_total_shifts()["worker3"] == 50
    assert (
        Rota.get_workers_total_shifts()["worker1"]
        + Rota.get_workers_total_shifts()["worker2"]
        == 50
    )

    worker3.pair = "1"

    Rota.build_and_solve()

    assert Rota.results.solver.status in ("warning", "error")

    worker3.pair = 0

    Rota.build_and_solve()
    Rota.export_rota_to_html("worker_pairs")

    assert Rota.results.solver.status == "ok"

    worker4 = Worker(name="worker4", site="group1", grade=1, fte=100, pair="1")
    Rota.add_worker(worker4)

    Rota.build_and_solve()
    Rota.export_rota_to_html("worker_pairs")

    assert Rota.results.solver.status == "ok"
    assert Rota.get_workers_total_shifts()["worker3"] == 50

def test_not_available_to_work():
    Rota = generate_basic_rota()

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=12.5,
            days=days,
            workers_required=1,
        ),
    )

    natws = generate_not_available_to_works(Rota.start_date, days=21)

    natws.append(
        NotAvailableToWork(date=Rota.rota_end_date - datetime.timedelta(days=1))
    )

    Rota.add_workers(
        [
            Worker(name="worker1", site="group1", grade=1),
            Worker(
                name="worker2", site="group1", grade=1, not_available_to_work=natws
            ),
        ]
    )

    Rota.build_and_solve()
    Rota.export_rota_to_html("test_not_available_to_work")

    assert Rota.results.solver.status == "ok"

    assert Rota.get_workers_total_shifts()["worker1"] == 35
    assert Rota.get_workers_total_shifts()["worker2"] == 35

    assert Rota.get_worker_shift_list_string(
        Rota.get_worker_by_name("worker2")
    ).startswith("-" * 21)
    assert Rota.get_worker_shift_list_string(
        Rota.get_worker_by_name("worker2")
    ).endswith("-")

def test_worker_ftes_2():
    Rota = generate_basic_rota(workers=2)

    Rota.add_workers(
        [
            Worker(
                name="worker3",
                site="group1",
                grade=1,
                fte=60,
            ),
            Worker(
                name="worker4",
                site="group1",
                grade=1,
                fte=40,
            ),
        ]
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:6],
            workers_required=1,
        ),
    )

    Rota.build_and_solve()

    Rota.export_rota_to_html("test_worker_ftes")

    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        match worker.name:
            case "worker1":
                assert worker.get_fte() == 100
                assert total_shifts[worker.name] == 20
            case "worker2":
                assert worker.get_fte() == 100
                assert total_shifts[worker.name] == 20
            case "worker3":
                assert worker.get_fte() == 60
                assert total_shifts[worker.name] == 12
            case "worker4":
                assert worker.get_fte() == 40
                assert total_shifts[worker.name] == 8

def test_worker_fte_shift_override():
    Rota = generate_basic_rota(workers=1)

    Rota.add_workers(
        [
            Worker(
                name="worker2",
                site="group1",
                grade=1,
                fte=100,
                shift_fte_overrides={"b": 50}
            ),
        ]
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:2],
            workers_required=1,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days[2:5],
            workers_required=1,
        ),
    )

    Rota.build_and_solve()

    Rota.export_rota_to_html("test_worker_fte_shift_override")

    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        match worker.name:
            case "worker1":
                assert worker.get_fte() == 100
                assert total_shifts[worker.name] == 30
                assert worker.get_fte("b") == 100
            case "worker2":
                assert worker.get_fte() == 100
                assert worker.get_fte("b") == 50
                assert total_shifts[worker.name] == 20

def test_worker_requirement_and_force_block():
    Rota = generate_basic_rota(workers=10, weeks_to_rota=10)

    wr = [
        WorkerRequirement(
            end_date=Rota.start_date + datetime.timedelta(weeks=5),
            number=2,
        ),
        WorkerRequirement(
            start_date=Rota.start_date + datetime.timedelta(weeks=5),
            number=3,
        ),
    ]

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:2],
            workers_required=1,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days[2:5],
            workers_required=1,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="c",
            length=12.5,
            days=days[:4],
            workers_required=wr,
            force_as_block=True,
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            ]
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="d",
            length=12.5,
            days=days[4:],
            workers_required=wr,
            force_as_block=True,
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            ]
        ),
    )

    Rota.build_and_solve()

    Rota.export_rota_to_html("test_worker_complex")

    assert Rota.results.solver.status == "ok"
    for worker in Rota.get_workers():
        shift_string = Rota.get_worker_shift_list_string(worker)
        assert shift_string.count("cccc") == shift_string.count("c") / 4
        assert shift_string.count("ddd") == shift_string.count("d") / 3

def test_split_shift():
    Rota = generate_basic_rota(workers=11, weeks_to_rota=22)

    Rota.set_rota_constraint("max_shifts_per_week", 2)
    Rota.set_rota_constraint("max_weekend_frequency", 3)
    Rota.set_rota_constraint("max_shifts_per_month", 4)

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:4],
            workers_required=1,
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            {"name": "max_shifts_per_week", "options": 1},
            ]
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=(days[4], days[6]),
            workers_required=1,
            assign_as_block=True,
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            ]
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="c",
            length=12.5,
            days=days[5],
            workers_required=1,
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            ]
        ),
    )

    Rota.build_and_solve()

    Rota.export_rota_to_html("test_split")
    assert Rota.results.solver.status == "ok"

def test_locums_basic():
    Rota = generate_basic_rota(workers=2, weeks_to_rota=6)

    la = []
    for date in Rota.get_date_range(end_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
        la.append(WorkRequests(date=date, shift="*"))
    la2 = []
    for date in Rota.get_date_range(start_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
        la2.append(WorkRequests(date=date, shift="*"))

    Rota.add_workers(
        [
            Worker(name="worker03", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1),
            Worker(name="worker04", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2),
            Worker(name="locum1", site="group1", locum=True),
            Worker(
                name="locum2", site="group1", locum=True
            ),
        ]
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:4],
            workers_required=1,
            
            constraint=[
            #{"name": "pre", "options": 1},
            #{"name": "post", "options": 1},
            #{"name": "max_shifts_per_week", "options": 1},
            ]
        ),
    )
    Rota.set_rota_constraint("distribute_locum_shifts", True)
    Rota.set_rota_constraint("maximum_allowed_shift_diff", 0)
    Rota.set_rota_constraint("maximum_allowed_locum_shifts_per_worker", 4)

    Rota.build_and_solve()

    Rota.export_rota_to_html("test_locums_basic", timestamp_filename=False)

    assert Rota.results.solver.status == "ok"

    for worker in Rota.get_workers():
        assert Rota.get_worker_shift_list(worker).count("a") == 4

def test_locums_nwds():
    Rota = generate_basic_rota(workers=2, weeks_to_rota=6)

    la = []
    for date in Rota.get_date_range(end_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
        la.append(WorkRequests(date=date, shift="*"))
    la2 = []
    for date in Rota.get_date_range(start_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
        la2.append(WorkRequests(date=date, shift="*"))

    Rota.add_workers(
        [
            Worker(name="worker03", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1, nwds=[NonWorkingDays(day="Tue")]),
            Worker(name="worker04", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2, nwds=[NonWorkingDays(day="Wed"),NonWorkingDays(day="Thu")]),
            Worker(name="locum1", site="group1", locum=True),
            Worker(
                name="locum2", site="group1", locum=True
            ),
        ]
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:4],
            workers_required=1,
            
            constraint=[
            #{"name": "pre", "options": 1},
            #{"name": "post", "options": 1},
            #{"name": "max_shifts_per_week", "options": 1},
            ]
        ),
    )
    Rota.set_rota_constraint("maximum_allowed_shift_diff", 0)
    Rota.set_rota_constraint("maximum_allowed_locum_shifts_per_worker", 4)

    Rota.build_and_solve()

    Rota.export_rota_to_html("test_locums_basic", timestamp_filename=False)

    assert Rota.results.solver.status == "ok"

    for worker in Rota.get_workers():
        print(  Rota.get_worker_shift_list(worker))
        assert Rota.get_worker_shift_list(worker).count("a") == 4

        if worker.name in ("worker03", "worker04"):
            assert Rota.get_worker_shift_list(worker, include_locums=True).count("a") > 6

def test_locums_no_availablity():
    Rota = generate_basic_rota(workers=2, weeks_to_rota=5)

    la = []
    la2 = []

    Rota.add_workers(
        [
            Worker(name="worker03", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1),
            Worker(name="worker04", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2),
            Worker(name="locum1", site="group1", locum=True),
            Worker(
                name="locum2", site="group1", locum=True
            ),
        ]
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:4],
            workers_required=1,
            
            constraint=[
            #{"name": "pre", "options": 1},
            #{"name": "post", "options": 1},
            #{"name": "max_shifts_per_week", "options": 1},
            ]
        ),
    )
    Rota.set_rota_constraint("distribute_locum_shifts", False)

    with pytest.raises(WarningTermination):
        Rota.build_and_solve()

def test_locums():
    Rota = generate_basic_rota(workers=7, weeks_to_rota=10)

    Rota.set_rota_constraint("max_shifts_per_week", 2)
    Rota.set_rota_constraint("max_weekend_frequency", 3)
    Rota.set_rota_constraint("max_shifts_per_month", 4)
    Rota.set_rota_constraint("distribute_locum_shifts", True)
    Rota.set_rota_constraint("maximum_allowed_shift_diff", 1)
    Rota.set_rota_constraint("maximum_allowed_locum_shifts_per_worker", 6)

    Rota.terminate_on_warning.remove("Worker/no valid shifts")

    la = []
    for date in Rota.get_date_range(end_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
        la.append(WorkRequests(date=date, shift="*"))
    la2 = []
    for date in Rota.get_date_range(start_date=Rota.rota_end_date-datetime.timedelta(weeks=3)):
        la2.append(WorkRequests(date=date, shift="*"))

    Rota.add_workers(
        [
            Worker(name="worker08", site="group1", locum=False, locum_availability=la, locum_max_shifts=40, locum_max_shifts_per_week=1),
            Worker(name="worker09", site="group1", locum=False, locum_availability=la2, locum_max_shifts=40, locum_max_shifts_per_week=2),
            Worker(name="locum1", site="group1", locum=True),
            Worker(
                name="locum2", site="group1", locum=True
            ),
        ]
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:4],
            workers_required=1,
            
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            {"name": "max_shifts_per_week", "options": 1},
            ]
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=(days[4], days[6]),
            workers_required=1,
            assign_as_block=True,
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            ]
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="c",
            length=12.5,
            days=days[5],
            workers_required=1,
            constraint=[
            {"name": "pre", "options": 2},
            {"name": "post", "options": 3},
            ]
        ),
    )

    Rota.build_and_solve()

    Rota.export_rota_to_html("test_locums", timestamp_filename=False)

    assert Rota.results.solver.status == "ok"

def test_worker_assign_as_block_preference():
    """Test that a worker's assign_as_block preference is respected over the global setting."""
    Rota = generate_basic_rota(workers=4, weeks_to_rota=12)

    # Worker 1 prefers to have shift 'a' assigned as a block, worker 2 does not
    workers = Rota.get_workers()
    worker1 = workers[0]
    worker2 = workers[1]
    worker3 = workers[2]
    worker4 = workers[3]

    # Set a strong preference for worker1 to have 'a' as a block
    worker1.assign_as_block_preferences = {"a":1}
    worker2.assign_as_block_preferences = {"a":-1}  # Worker 2 does not want 'a' as a block

    # Add a shift 'a' that is NOT globally set as assign_as_block
    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            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_assign_as_block_preference", folder="tests")

    # Worker 1 should have their 'a' shifts grouped together (block), worker 2 may not
    shift_list_1 = Rota.get_worker_shift_list_string(worker1)
    assert shift_list_1.count("aaaa") == 3, "Worker 1 should have all 'a' shifts grouped together (block)"

    shift_list_2 = Rota.get_worker_shift_list_string(worker2)
    assert shift_list_2.count("a") == 12, "Worker 2 should have at 12 'a' shift, but may not be grouped together"

    shift_list_3 = Rota.get_worker_shift_list_string(worker3)
    assert shift_list_3.count("a") == 12, "Worker 3 should have at least 12 'a' shifts, but may not be grouped together"

    shift_list_4 = Rota.get_worker_shift_list_string(worker4)
    assert shift_list_4.count("a") == 12, "Worker 4 should have at least 12 'a' shifts, but may not be grouped together"

    # Modify worker1's preference to prevent block 'a' shifts
    worker1.assign_as_block_preferences = {"a":-1}
    Rota.build_and_solve(options={"ratio": 0.0})
    shift_list_1 = Rota.get_worker_shift_list_string(worker1)
    assert shift_list_1.count("aaaa") == 0, "Worker 1 should have all 'a' shifts grouped together (block)"

    worker1.assign_as_block_preferences = {"a":0}

    a_shift = Rota.get_shift_by_name("a")
    a_shift.assign_as_block = True  # Set the shift 'a' to be assigned as a block globally
    Rota.build_and_solve(options={"ratio": 0.0})

    for worker in Rota.get_workers():
        shift_list = Rota.get_worker_shift_list_string(worker)
        # All workers should have 'a' shifts grouped together (block)
        assert shift_list.count("aaaa") == 3, f"Worker {worker.name} should have all 'a' shifts grouped together (block)"

def test_worker_assign_as_block_preference_multiple_shifts():
    """Test that a worker's assign_as_block preference is respected over the global setting."""
    Rota = generate_basic_rota(workers=4, weeks_to_rota=12)

    # Worker 1 prefers to have shift 'a' assigned as a block, worker 2 does not
    workers = Rota.get_workers()
    worker1 = workers[0]
    worker2 = workers[1]
    worker3 = workers[2]
    worker4 = workers[3]

    # Set a strong preference for worker1 to have 'a' as a block
    worker1.assign_as_block_preferences = {"a":100}
    worker2.assign_as_block_preferences = {"b":-10000}  # Worker 2 does not want 'a' as a block

    # Add a shift 'a' that is NOT globally set as assign_as_block
    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})

    shift_list_1 = Rota.get_worker_shift_list_string(worker1)
    assert shift_list_1.count("aaaa") == 3, "Worker 1 should have all 'a' shifts grouped together (block)"
    Rota.export_rota_to_html("test_worker_assign_as_block_preference", folder="tests")

def test_worker_assign_as_block_preference_multiple_sites():
    """Test that a worker's assign_as_block preference is respected over the global setting."""
    Rota = generate_basic_rota(workers=6, weeks_to_rota=12)

    # Worker 1 prefers to have shift 'a' assigned as a block, worker 2 does not
    workers = Rota.get_workers()
    worker1 = workers[0]
    worker2 = workers[1]
    worker3 = workers[2]
    worker4 = workers[3]
    worker5 = workers[4]
    worker6 = workers[5]

    # Set a strong preference for worker1 to have 'a' as a block
    worker1.assign_as_block_preferences = {"a":1}

    worker4.site = "group2"  
    worker5.site = "group2"  
    worker6.site = "group2" 
    worker4.assign_as_block_preferences = {"b":-10000}  # Worker 2 does not want 'a' as a block

    # Add a shift 'a' that is NOT globally set as assign_as_block
    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=("group2",),
            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_assign_as_block_preference", folder="tests")

def test_worker_hard_day_dependency():
    Rota = generate_basic_rota(workers=6, weeks_to_rota=16)

    workers = Rota.get_workers()

    worker1 = workers[0]
    worker2 = workers[1]
    worker3 = workers[2]

    worker1.add_hard_day_dependency("Mon", "Tue")
    worker2.add_hard_day_dependency("Mon", "Wed")
    worker3.add_hard_day_dependency("Mon", "Thu")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days,
            workers_required=2,
            assign_as_block=False,  # global setting is off
        ),
    )

    Rota.build_and_solve(options={"ratio": 0.0})
    Rota.export_rota_to_html("test_worker_hard_day_depend", folder="tests")

    for worker in Rota.get_workers():
        shift_list = Rota.get_worker_shift_list_string(worker)

        if worker.name == "worker1":
            # Check that every time 'a' is assigned on Monday, it is also assigned on Tuesday
            for week in range(12):
                mon_idx = week * 7 + days.index("Mon")
                tue_idx = week * 7 + days.index("Tue")
                if shift_list[mon_idx] == "a":
                    assert shift_list[tue_idx] == "a", f"Worker1 assigned 'a' on Mon (week {week}) but not on Tue"

        if worker.name == "worker2":
            # Check that every time 'a' is assigned on Monday, it is also assigned on Wednesday
            for week in range(12):
                mon_idx = week * 7 + days.index("Mon")
                wed_idx = week * 7 + days.index("Wed")
                if shift_list[mon_idx] == "a":
                    assert shift_list[wed_idx] == "a", f"Worker2 assigned 'a' on Mon (week {week}) but not on Wed"

        if worker.name == "worker3":
            # Check that every time 'a' is assigned on Monday, it is also assigned on Thursday
            for week in range(12):
                mon_idx = week * 7 + days.index("Mon")
                thu_idx = week * 7 + days.index("Thu")
                if shift_list[mon_idx] == "a":
                    assert shift_list[thu_idx] == "a", f"Worker3 assigned 'a' on Mon (week {week}) but not on Thu"

                
def test_worker_hard_day_dependency_weekend():
    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)

    workers = Rota.get_workers()

    worker1 = workers[0]
    worker2 = workers[1]
    worker3 = workers[2]
    worker4 = workers[3]

    worker1.add_hard_day_dependency("Fri", "Sun")
    worker2.add_hard_day_dependency("Fri", "Sun")
    worker3.add_hard_day_dependency("Fri", "Sun")
    worker4.add_hard_day_dependency("Fri", "Sun")

    worker1.add_hard_day_exclusion("Fri", "Sat")
    worker2.add_hard_day_exclusion("Fri", "Sat")
    worker3.add_hard_day_exclusion("Fri", "Sat")
    worker4.add_hard_day_exclusion("Fri", "Sat")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            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_hard_day_depend", folder="tests")

    for worker in Rota.get_workers():
        shift_list = Rota.get_worker_shift_list_string(worker)

        # Check that every time 'a' is assigned on Monday, it is also assigned on Tuesday
        for week in range(16):
            fri_idx = week * 7 + days.index("Fri")
            sat_idx = week * 7 + days.index("Sat")
            sun_idx = week * 7 + days.index("Sun")
            if shift_list[fri_idx] == "a":
                assert shift_list[sun_idx] == "a", f"Worker1 assigned 'a' on Fri (week {week}) but not on Sun"

                assert shift_list[sat_idx] != "a", f"Worker1 assigned 'a' on Fri (week {week}) but also on Sat, which should not happen due to hard day exclusion"

            
def test_worker_hard_day_dependency_weekend2():
    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)

    workers = Rota.get_workers()

    worker1 = workers[0]
    worker2 = workers[1]
    worker3 = workers[2]
    worker4 = workers[3]

    worker1.add_hard_day_dependency("Fri", "Sun")
    worker2.add_hard_day_dependency("Fri", "Sun")
    worker3.add_hard_day_dependency("Fri", "Sun")
    worker4.add_hard_day_dependency("Fri", "Sun")
    worker4.add_hard_day_dependency("Fri", "Sat")

    worker1.add_hard_day_exclusion("Fri", "Sat")
    worker2.add_hard_day_exclusion("Fri", "Sat")
    worker3.add_hard_day_exclusion("Fri", "Sat")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            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_hard_day_depend", folder="tests")

    for worker in Rota.get_workers():
        shift_list = Rota.get_worker_shift_list_string(worker)

        # Check that every time 'a' is assigned on Monday, it is also assigned on Tuesday
        for week in range(16):
            fri_idx = week * 7 + days.index("Fri")
            sat_idx = week * 7 + days.index("Sat")
            sun_idx = week * 7 + days.index("Sun")
            if worker != worker4 and shift_list[fri_idx] == "a":
                assert shift_list[sun_idx] == "a", f"Worker1 assigned 'a' on Fri (week {week}) but not on Sun"

                assert shift_list[sat_idx] != "a", f"Worker1 assigned 'a' on Fri (week {week}) but also on Sat, which should not happen due to hard day exclusion"

            if worker == worker4 and shift_list[fri_idx] == "a":
                assert shift_list[sun_idx] == "a", f"Worker4 assigned 'a' on Fri (week {week}) but not on Sun"
                assert shift_list[sat_idx] == "a", f"Worker4 assigned 'a' on Fri (week {week}) but not on Sun"


def test_worker_force_assign_with_invalid_shift():
    Rota = generate_basic_rota(workers=4, weeks_to_rota=2)
    workers = Rota.get_workers()
    for worker in workers:
        worker.add_force_assign_with("a", "c")

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

    with pytest.raises(InvalidShift):
        # This should raise an InvalidShift error because 'c' does not exist
        Rota.build_and_solve(options={"ratio": 0.0})
    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="c",
            length=8,
            days=days[:4],
            workers_required=1,
            assign_as_block=False,  # global setting is off
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
            

def test_worker_force_assign_with():
    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)

    workers = Rota.get_workers()

    for worker in workers:
        worker.add_force_assign_with("a", "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_with", folder="tests")

    for worker in Rota.get_workers():
        shift_list = Rota.get_worker_shift_list(worker, search_multiple_assignments=True)

        for shift in shift_list:
            if shift == "a":
                assert "b" in shift_list, f"Worker {worker.name} has 'a' but not 'b' assigned with it"
            elif shift == "b":
                assert "a" in shift_list, f"Worker {worker.name} has 'b' but not 'a' assigned with it"

        assert shift_list.count({"a", "b"}) == 16

def test_worker_force_assign_with_and_hard_days():
    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)

    workers = Rota.get_workers()

    for worker in workers:
        worker.add_force_assign_with("a", "b")

        #worker.add_hard_day_dependency("Sat", "Sun")
        worker.add_hard_day_dependency("Mon", "Tue")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days[:2],
            workers_required=1,
            assign_as_block=False,  # global setting is off
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[:2],
            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_with_hard_days", folder="tests")

    assert Rota.results.solver.status == "ok"

def test_worker_force_assign_with_and_hard_days_weekend():
    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)

    workers = Rota.get_workers()

    for worker in workers:
        worker.add_force_assign_with("a", "b")

        #worker.add_hard_day_dependency("Sat", "Sun")
        worker.add_hard_day_dependency("Sat", "Sun")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days[5:],
            workers_required=1,
            assign_as_block=False,  # global setting is off
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[5:],
            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_with_hard_days", folder="tests")

    assert Rota.results.solver.status == "ok"


def test_workers_double_shifts_on_mon_tue():
    Rota = generate_basic_rota(workers=2, weeks_to_rota=16)

    workers = Rota.get_workers()

    for worker in workers:
        worker.add_force_assign_with("a", "b")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days[:2],
            workers_required=2,
            assign_as_block=False,  # global setting is off
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[:2],
            workers_required=2,
            assign_as_block=False,  # global setting is off
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")

    assert Rota.results.solver.status == "ok"

def test_workers_double_shifts_on_weekends():
    Rota = generate_basic_rota(workers=2, weeks_to_rota=16)

    workers = Rota.get_workers()

    for worker in workers:
        worker.add_force_assign_with("a", "b")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days[5:],
            workers_required=2,
            assign_as_block=False,  # global setting is off
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[5:],
            workers_required=2,
            assign_as_block=False,  # global setting is off
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")

    assert Rota.results.solver.status == "ok"


def test_max_shifts_per_week_by_shift_name_basic():
    Rota = RotaBuilder(weeks_to_rota=2, start_date=datetime.date(2025, 6, 16))
    w = Worker(name="A01", site="site1", grade=1)
    w.max_shifts_per_week_by_shift_name = {"a": 1}
    Rota.add_worker(w)
    Rota.add_shifts(
        SingleShift(
            sites=("site1",),
            name="a",
            length=8,
            days=days[:5],
            workers_required=1,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    # Worker should not be assigned more than 1 "a" shift per week
    for week in Rota.weeks:
        count = sum(
            Rota.model.works[w.id, week, day, "a"].value > 0.5
            for day in days[:5]
        )
        assert count <= 1

def test_max_shifts_per_week_by_shift_name_multiple_weeks():
    Rota =generate_basic_rota(workers=4, weeks_to_rota=2)
    for worker in Rota.get_workers():
        worker.set_max_shifts_per_week("a", 1)
    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days[:5],
            workers_required=1,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    assert Rota.results.solver.status == "error"

    Rota.add_worker(
        Worker(name="extraworker", site="group1", grade=1, max_shifts_per_week_by_shift_name={"a": 1})
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    assert Rota.results.solver.status == "ok"

    Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
    # Worker should not be assigned more than 2 "a" shifts per week
    for worker in Rota.get_workers():
        for week in Rota.weeks:
            count = sum(
                Rota.model.works[worker.id, week, day, "a"].value > 0.5
                for day in days[:5]
            )
            assert count <= 2


def test_max_shifts_per_week_by_shift_name_invalid_shift():
    Rota = RotaBuilder(weeks_to_rota=1, start_date=datetime.date(2025, 6, 16))
    w = Worker(name="A03", site="site1", grade=1)
    # "b" does not exist
    w.max_shifts_per_week_by_shift_name = {"b": 1}
    Rota.add_worker(w)
    Rota.add_shifts(
        SingleShift(
            sites=("site1",),
            name="a",
            length=8,
            days=days[:5],
            workers_required=1,
        ),
    )
    with pytest.raises(InvalidShift):
        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"