proc-rota/test/test_rota.py

from copy import deepcopy
from re import S
from black import main
import pytest
from pytest import approx
from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, days

import datetime

from rota.workers import Worker


@pytest.fixture
def demo_rota():
    weeks_to_rota = 10
    start_date = datetime.date(2022, 3, 7)
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
        balance_offset_modifier=1,
    )
    Rota.constraint_options["max_shifts_per_week"] = 5
    Rota.constraint_options["max_shifts_per_month"] = 31
    Rota.constraint_options["constrain_time_off_after_nights"] = False
    Rota.constraint_options["balance_nights_across_site"] = False
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.constraint_options["balance_shifts_over_workers"] = True
    Rota.constraint_options["balance_nights"] = False
    Rota.constraint_options["minimise_shift_diffs"] = False
    Rota.constraint_options["balance_blocks"] = False
    Rota.constraint_options["balance_bank_holidays"] = False

    # Add a few workers
    workers = [
        Worker(name="worker1", site="group1", grade=1),
        Worker(name="worker2", site="group1", grade=1),
        Worker(name="worker3", site="group1", grade=1),
        Worker(name="worker4", site="group1", grade=1),
        Worker(name="worker5", site="group2", grade=1, fte=60),
        Worker(name="worker6", site="group2", grade=1, fte=40),
    ]
    Rota.add_workers(workers)

    # Add a weekday and weekend shift
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="weekday",
            length=12.5,
            days=days[:5],
            balance_offset=1,
            workers_required=1,
        ),
        SingleShift(
            sites=("group1",),
            name="weekend",
            length=12.5,
            days=days[5:],
            balance_offset=4,
        ),
    )

    Rota.build_shifts()
    Rota.build_workers()
    Rota.build_model()

    solver_options = {"ratio": 0.1, "seconds": 1000, "threads": 10}
    Rota.solve_model(options=solver_options)
    Rota.export_rota_to_html("test2")
    return Rota, start_date, weeks_to_rota

def test_optimal_solution(demo_rota):
    Rota, _, _ = demo_rota
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_start_date(demo_rota):
    Rota, start_date, _ = demo_rota
    assert Rota.start_date == start_date

def test_end_date(demo_rota):
    Rota, start_date, weeks_to_rota = demo_rota
    assert Rota.rota_end_date == start_date + datetime.timedelta(weeks=weeks_to_rota)

def test_rota_workers_number(demo_rota):
    Rota, _, _ = demo_rota
    assert len(Rota.workers) == 6

def test_rota_workers_total_fte(demo_rota):
    Rota, _, _ = demo_rota
    assert Rota.get_workers_total_fte() == 500

def test_weekday_workers(demo_rota):
    Rota, _, _ = demo_rota
    weekday_shift_workers = Rota.get_workers_for_shift(Rota.get_shift_by_name("weekday"))
    for worker in weekday_shift_workers:
        assert worker.proportion_rota_to_work == 1
    assert len(weekday_shift_workers) == 6

def test_weekend_workers(demo_rota):
    Rota, _, _ = demo_rota
    weekend_shift_workers = Rota.get_workers_for_shift(Rota.get_shift_by_name("weekend"))
    for worker in weekend_shift_workers:
        assert worker.proportion_rota_to_work == 1
    assert len(weekend_shift_workers) == 4

def test_worker_shift_targets(demo_rota):
    Rota, _, _ = demo_rota
    for worker in Rota.workers:
        match worker.site:
            case "group1":
                assert worker.proportion_rota_to_work == 1
                assert worker.get_shift_targets() == {"weekday": 10, "weekend": 5}
            case "group2":
                assert worker.proportion_rota_to_work == 1
                match worker.fte:
                    case 60:
                        assert worker.get_shift_targets() == {"weekday": 6, "weekend": 0}
                    case 40:
                        assert worker.get_shift_targets() == {"weekday": 4, "weekend": 0}
                    case _:
                        assert 1 == 0
            case _:
                assert 1 == 0

def test_worker_shift_targets2(demo_rota):
    Rota, _, _ = demo_rota
    shift_number = Rota.get_shift_summary_dict()
    for worker in Rota.workers:
        match worker.site:
            case "group1":
                assert shift_number[worker.name]["weekday"] == 10
                assert shift_number[worker.name]["weekend"] == 5
            case "group2":
                assert shift_number[worker.name]["weekend"] == 0
                match worker.fte:
                    case 60:
                        assert shift_number[worker.name]["weekday"] == 6
                    case 40:
                        assert shift_number[worker.name]["weekday"] == 4
                    case _:
                        assert 1 == 0
            case _:
                assert 1 == 0

def test_rota_sites(demo_rota):
    Rota, _, _ = demo_rota
    assert Rota.sites == {"group1", "group2"}





 





# class TestNoWorkerRota:
#    weeks_to_rota = 10
#    start_date = datetime.date(2022, 3, 7)
#
#    Rota = RotaBuilder(
#        start_date,
#        weeks_to_rota=weeks_to_rota,
#    )
#
#    def test_no_workers(self):
#        with pytest.raises(NoWorkers):
#            self.Rota.build_workers()