proc-rota/test/rota_constraints/test_avoid_workers_by_week.py

import datetime
import pytest
from rota_generator.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days

from rota_generator.workers import Worker

import itertools


def generate_basic_rota(weeks_to_rota=9):
    start_date = datetime.date(2022, 3, 7)

    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

    # Add a few workers
    Rota.add_workers(
        [
            Worker(name="worker1", site="group1", grade=1),
            Worker(name="worker2", site="group1", grade=2),
            Worker(name="worker3", site="group1", grade=3),
        ]
    )

    return Rota


def date_generator(from_date, days):
    n = 0
    while True:
        yield from_date

        n = n + 1
        if n >= days:
            break
        from_date = from_date + datetime.timedelta(days=1)


class TestWorkerRequests:
    def test_avoid_worker_constraint(self):
        Rota = generate_basic_rota()


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

        Rota.add_worker_name_constraint_by_week(["worker1"], [1,2,3], ["a"])
        Rota.add_worker_name_constraint_by_week(["worker2"], [4,5,6], ["a"])
        Rota.add_worker_name_constraint_by_week(["worker3"], [7,8,9], ["a"])

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("avoid_grade_constraint")

        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)

            shift_string = Rota.get_worker_shift_list_string(worker)

            match worker_name:
                case "worker1":
                    assert shift_string.startswith("-"*21)
                case "worker2":
                    assert shift_string.startswith(21*"a"+"-"*21)
                case "worker3":
                    assert shift_string.endswith("-"*21)


    def test_avoid_worker_constraint_invalid_shift(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="a", length= 12.5, days=days,
                workers_required=1,
                force_as_block=False,
            ),
        )
        with pytest.raises(InvalidShift):

            Rota.add_worker_name_constraint_by_week(["worker1"], [1,2,3], ["b"])
            Rota.build_and_solve(options={"ratio": 0.000})


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

        Rota.add_worker_name_constraint_by_week([1], [1,2,3], ["b"])
        Rota.build_and_solve(options={"ratio": 0.000})
        assert Rota.results.solver.status == "ok"


# TODO: implement check for invalid workers
##    def test_avoid_grade_constraint_invalid_worker(self):
##        Rota = generate_basic_rota()
##
##        Rota.add_shifts(
##            SingleShift(
##                sites=("group1",), name="a", length= 12.5, days=days,
##                workers_required=2,
##                force_as_block=False,
##            ),
##        )
##        with pytest.raises(ValueError):
##
##            Rota.add_worker_name_constraint_by_week([4], [1,2,3], ["a"])
##            Rota.build_and_solve(options={"ratio": 0.000})
##
##        Rota.add_worker(
##                Worker(name="worker4", site="group1", grade=4),
##        )
##
##
##        Rota.add_worker_name_constraint_by_week([4], [1,2,3], ["a"])
##        Rota.build_and_solve(options={"ratio": 0.000})
##        assert Rota.results.solver.status == "ok"

    def test_avoid_worker_constraint_multiple_workers(self):
        Rota = generate_basic_rota()

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

        Rota.add_worker_name_constraint_by_week(["worker1", "worker2"], [1,2,3], ["a"])
        Rota.build_and_solve(options={"ratio": 0.000})
        assert Rota.results.solver.status == "ok"

        Rota.add_worker_name_constraint_by_week(["worker3"], [1,2,3], ["a"])
        Rota.build_and_solve(options={"ratio": 0.000})
        assert Rota.results.solver.status in ("warning", "error")

    def test_avoid_worker_constraint_multiple_shifts(self):
        Rota = generate_basic_rota()

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

        Rota.add_worker_name_constraint_by_week(["worker1"], [1,2,3], ["a", "b"])
        Rota.add_worker_name_constraint_by_week(["worker2"], [4,5,6], ["a", "b"])
        Rota.add_worker_name_constraint_by_week(["worker3"], [7,8,9], ["a", "b"])

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("avoid_grade_constraint")

        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)

            shift_string = Rota.get_worker_shift_list_string(worker)

            match worker_name:
                case "worker1":
                    assert shift_string.startswith("-"*21)
                case "worker2":
                    assert "-"*21 in shift_string # hacky
                case "worker3":
                    assert shift_string.endswith("-"*21)

        Rota.add_worker_name_constraint_by_week(["worker2"], [7,8,9], ["a"])
        Rota.build_and_solve(options={"ratio": 0.000})
        assert Rota.results.solver.status == "ok"

        Rota.add_worker_name_constraint_by_week(["worker1"], [7,8,9], ["a"])
        Rota.build_and_solve(options={"ratio": 0.000})
        assert Rota.results.solver.status in ("warning", "error")