proc-rota/test_rota.py

from black import main
import pytest
from shifts import NoWorkers, RotaBuilder, RotaResults, SingleShift, days

import datetime

from workers import Worker


class TestDemoRota:
    # Set up 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,
        max_weekend_frequency=1,
    )

    Rota.constraint_options["max_shifts_per_week"] = 5
    Rota.constraint_options["max_shifts_per_month"] = 31
    Rota.constraint_options["ensure_derriford_reg_for_nights"] = False
    Rota.constraint_options["limit_to_1_st2_on_nights"] = False
    Rota.constraint_options["ensure_1_st4_plus_on_nights"] = False
    Rota.constraint_options["constrain_time_off_after_nights"] = False
    Rota.constraint_options["balance_nights_across_site"] = False
    Rota.constraint_options["balance_shifts"] = True
    Rota.constraint_options["balance_nights"] = False
    Rota.constraint_options["minimise_shift_diffs"] = True
    #Rota.constraint_options["balance_weekends"] = False

    # Add a few workers
    worker1 = Worker(Rota, "worker1", "group1", 1)
    worker2 = Worker(Rota, "worker2", "group1", 1)
    worker3 = Worker(Rota, "worker3", "group1", 1)
    worker4 = Worker(Rota, "worker4", "group1", 1)
    worker5 = Worker(Rota, "worker5", "group2", 1, fte=60)
    worker6 = Worker(Rota, "worker6", "group2", 1, fte=40)

    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6))

    # Add a weekday and weekend shift
    Rota.add_shifts(
        SingleShift(
            ("group1", "group2"),
            "weekday",
            12.5,
            days[:5],
            balance_offset=10,
            workers_required=1,
            # We use a different balance weighting for each shift as otherwise
            # the shifts can be treated equally (and no target diff)
            balance_weighting=0.5, 
        ),
        SingleShift(("group1",), "weekend", 12.5, days[5:], balance_offset=40,
        #balance_weighting=0.5
        ),
    )

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

    print(Rota.get_worker_details())

    for w in Rota.get_workers():
        print(w)
        print(w.get_shift_targets())

    solver_options = {"ratio": 0, "seconds": 1000, "threads": 10}

    Rota.solve_model(options=solver_options)

    # print(Rota.get_worker_details())

    # optimizer = SolverFactory('cbc')
    # result = optimizer.solve(prob,tee=True)
    # result.Solver.Status = SolverStatus.warning
    # prob.solutions.load_from(result)

    ResultsHolder = RotaResults(Rota)


    worker_timetable_brief = ResultsHolder.get_worker_timetable_brief(
        show_prefs=False, show_unavailable=False
    )

    print(worker_timetable_brief)
    ResultsHolder.export_rota_to_html("test")

    def test_start_date(self):
        assert self.Rota.start_date == self.start_date

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

    def test_rota_workers_number(self):
        assert len(self.Rota.workers) == 6

    def test_rota_workers_total_fte(self):
        assert self.Rota.get_workers_total_fte() == 500

    def test_weekday_workers(self):
        weekday_shift_workers = self.Rota.get_workers_for_shift(
            self.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(self):
        weekend_shift_workers = self.Rota.get_workers_for_shift(
            self.Rota.get_shift_by_name("weekend")
        )

        for worker in weekend_shift_workers:
            assert worker.proportion_rota_to_work == 1

            # assert worker.get_shift_targets() == {}

        assert len(weekend_shift_workers) == 4

    def test_worker_shift_targets(self):
        for worker in self.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_rota_sites(self):
        assert self.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()


if __name__ == "__main__":
    t = TestDemoRota()