further tests and fixes

rota/shifts.py

@@ -38,7 +38,7 @@ for bank_holiday in bank_holidays.get_holidays(division="england-and-wales"):
     bank_holiday_map[bank_holiday["date"]] = bank_holiday["title"]
 
 SHIFT_BOUNDS = {
-    "bank_holiday": (0, 4),
+    "bank_holiday": (0, 9),
     "shift_count": (0, 400),
     "night_shift_count": (0, 19),
     "weekend_count": (0, 60),
@@ -76,7 +76,7 @@ class SingleShift(object):
         length: float,
         shift_days: Sequence[str],
         # balance_by_site=True,
-        balance_offset: float = 2,  # this could be generated dynamically
+        balance_offset: float | None = None,
         balance_weighting: float = 1,
         workers_required: float = 1,
         rota_on_nwds: bool = False,
@@ -95,7 +95,11 @@ class SingleShift(object):
         # balance_offset defines the max difference in allocated shifts
         # versus target shifts (hard constraint)
         # if 0 exactly equal shifts must be assigend (unlikely to be possible)
-        self.balance_offset = balance_offset
+
+        if balance_offset is None:
+            self.balance_offset = len(shift_days)
+        else:
+            self.balance_offset = balance_offset
 
         # weight the shift for balancing, default is 1
         self.balance_weighting = balance_weighting
@@ -615,6 +619,12 @@ class RotaBuilder(object):
             initialize=availability_init,
         )
 
+        # Validate shifts are valid
+        work_request_sets = set([i[3] for i in self.work_requests])
+        invalid_shifts = work_request_sets - set(self.get_shift_names())
+        if invalid_shifts:
+            raise InvalidShift(f"Invalid worker shift request [shift(s): ${invalid_shifts}]")
+
         def work_request_init(model, wid, week, day, shift):
             if (wid, week, day, shift) in self.work_requests:
                 self.work_requests_map[(wid, week, day)] = shift
@@ -2016,6 +2026,7 @@ class RotaBuilder(object):
         """
         Process the added shifts
         """
+
         self.shifts_by_name = {}
         self.shift_names = []  # type: List[ShiftName]
 
@@ -2026,6 +2037,19 @@ class RotaBuilder(object):
         self.week_day_shift_product = []
         self.week_day_shiftclass_product = []
 
+        for s in self.shifts:
+            if s.name in self.shift_names:
+                raise InvalidShift(f"Duplicate shift: {s.name}")
+
+            self.shifts_by_name[s.name] = s
+            self.shift_names.append(s.name)
+
+            # for day in s.shift_days:
+            #    self.week_day_shifts_dict[(week, day)].add(s.name)
+
+            for site in s.site:
+                self.sites.add(site)
+
         self.shift_counts = defaultdict(int)
         for week, day in self.weeks_days_product:
             self.week_day_shifts_dict[(week, day)] = set()
@@ -2043,15 +2067,6 @@ class RotaBuilder(object):
                         self.week_day_shifts_dict[(week, day)].add(s.name)
                         self.shift_counts[s] = self.shift_counts[s] + 1
         # print(self.week_day_shift_product)
-        for s in self.shifts:
-            self.shifts_by_name[s.name] = s
-            self.shift_names.append(s.name)
-
-            # for day in s.shift_days:
-            #    self.week_day_shifts_dict[(week, day)].add(s.name)
-
-            for site in s.site:
-                self.sites.add(site)
 
     #        # Todo replace with week_day.....
     #        self.day_shift_product = []
@@ -2762,11 +2777,14 @@ class RotaBuilder(object):
 
 class NoActiveSites(Exception):
     """Raised when there are no active sites"""
-
     pass
 
 
 class NoWorkers(Exception):
     """Raised when there are no active sites"""
-
     pass
+
+
+class InvalidShift(Exception):
+    """Raised when there are no active sites"""
+    pass

rota/workers.py

@@ -1,6 +1,6 @@
 import datetime
 from collections import defaultdict
-from typing import List
+from typing import Iterable, List
 
 # from .shifts import RotaBuilder, days, sites
 import uuid
@@ -24,10 +24,10 @@ class Worker:
         nwd: list[str] | list[str, datetime.datetime, datetime.datetime] = [],
         start_date: datetime.datetime | None = None,
         end_date: datetime.datetime | None = None,
-        oop:list[(datetime.datetime, datetime.datetime, str)] | None = None,
+        oop: list[(datetime.datetime, datetime.datetime, str)] | None = None,
-        not_available_to_work=None,
+        not_available_to_work: list[(datetime.datetime, str)] | None = None,
-        pref_not_to_work=None,
+        pref_not_to_work: list[datetime.datetime] | None = None,
-        work_requests=None,
+        work_requests: list[(datetime.datetime, str)] | None = None,
         remote_site: str = "plymouth",  # We set a default proc_site
         previous_shifts: dict = {},
         shift_balance_extra: dict = {},
@@ -58,8 +58,11 @@ class Worker:
 
         self.shift_target_number = defaultdict(int)
 
-        #days_to_work = Rota.rota_days_length
+        # days_to_work = Rota.rota_days_length
 
+        if start_date is not None and end_date is not None:
+            if start_date >= end_date:
+                raise ValueError("End date must be after start date")
 
         # if no start date default to the start of the rota
         if start_date is None:
@@ -71,18 +74,12 @@ class Worker:
         if end_date is None:
             self.end_date = Rota.rota_end_date
         else:
-            if isinstance(end_date, datetime.date):
+            self.end_date = end_date
-                self.end_date = end_date
+            # self.end_date = datetime.datetime.strptime(end_date, "%d/%m/%y").date()
-            else:
-                self.end_date = datetime.datetime.strptime(end_date, "%d/%m/%y").date()
 
             if self.end_date > Rota.rota_end_date:
                 self.end_date = Rota.rota_end_date
 
-            
-        #if self.start_date >= self.end_date:
-        #    raise ValueError("End date must be after start date")
-
         days_to_work = (self.end_date - self.start_date).days
 
         for week, day in Rota.weeks_days_product:
@@ -94,42 +91,28 @@ class Worker:
                     (self.id, week, day)
                 ] = f"START DATE: {self.start_date}"
 
-            if date >= self.end_date: 
+            if date >= self.end_date:
                 Rota.unavailable_to_work.add((self.id, week, day))
                 Rota.unavailable_to_work_reason[
                     (self.id, week, day)
                 ] = f"END DATE: {self.end_date}"
 
-        #if self.start_date > Rota.start_date:
-        #    # add unavalabilities
-        #    for weeks_days in Rota.weeks_days_product[:days_to_work]:
-        #        week, day = weeks_days
-        #        Rota.unavailable_to_work.add((self.id, week, day))
-        #        Rota.unavailable_to_work_reason[
-        #            (self.id, week, day)
-        #        ] = f"START DATE: {self.start_date}"
-
-        #if self.end_date < Rota.rota_end_date:
-        #    # add unavalabilities
-        #    for weeks_days in Rota.weeks_days_product[days_to_work:]:
-        #        week, day = weeks_days
-        #        Rota.unavailable_to_work.add((self.id, week, day))
-        #        Rota.unavailable_to_work_reason[
-        #            (self.id, week, day)
-        #        ] = f"END DATE: {self.end_date}"
-
         if oop is not None:
             for start_oop, end_oop, oop_name in oop:
-                #start_oop, end_oop = oop
+                # start_oop, end_oop = oop
                 if isinstance(start_oop, datetime.date):
                     start_oop_date = start_oop
                 else:
-                    start_oop_date = datetime.datetime.strptime(start_oop, "%d/%m/%y").date()
+                    start_oop_date = datetime.datetime.strptime(
+                        start_oop, "%d/%m/%y"
+                    ).date()
 
                 if isinstance(end_oop, datetime.date):
                     end_oop_date = end_oop
                 else:
-                    end_oop_date = datetime.datetime.strptime(end_oop, "%d/%m/%y").date()
+                    end_oop_date = datetime.datetime.strptime(
+                        end_oop, "%d/%m/%y"
+                    ).date()
 
                 if start_oop_date >= end_oop_date:
                     raise ValueError("End OOP date must be after start date")
@@ -163,7 +146,7 @@ class Worker:
             # loop throught dates converting to week / day combination
             for date in pref_not_to_work:
                 days_from_start = (
-                    datetime.datetime.strptime(date, "%d/%m/%y").date()
+                    date
                     - Rota.start_date
                 ).days
                 week = days_from_start // 7 + 1
@@ -180,7 +163,7 @@ class Worker:
             # loop throught dates converting to week / day combination
             for date, reason in not_available_to_work:
                 days_from_start = (
-                    datetime.datetime.strptime(date, "%d/%m/%y").date()
+                    date
                     - Rota.start_date
                 ).days
                 week = days_from_start // 7 + 1
@@ -191,7 +174,7 @@ class Worker:
         if work_requests is not None:
             for date, shift in work_requests:
                 days_from_start = (
-                    datetime.datetime.strptime(date, "%d/%m/%y").date()
+                    date
                     - Rota.start_date
                 ).days
                 week = days_from_start // 7 + 1

test/test_general_balancing.py

@@ -0,0 +1,280 @@
+import datetime
+import pytest
+from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
+
+from rota.workers import Worker
+
+
+def generate_basic_rota(weeks_to_rota=10):
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.constraint_options["balance_weekends"] = True
+
+    # Add a few workers
+    Rota.add_workers([
+        Worker(Rota, "worker1", "group1", 1),
+        Worker(Rota, "worker2", "group1", 1),
+        Worker(Rota, "worker3", "group2", 1),
+        #Worker(Rota, "worker4", "group2", 1),
+        #Worker(Rota, "worker5", "group2", 1),
+        #Worker(Rota, "worker6", "group2", 1, fte=50),
+    ])
+
+    return Rota
+
+
+class TestBalancing:
+    def test_basic_balancing(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                balance_offset=2,
+                workers_required=1,
+                force_as_block=False
+            ),
+            SingleShift(
+                ("group1", "group2"),
+                "b",
+                12.5,
+                days,
+                balance_offset=2,
+                workers_required=1,
+                force_as_block=False
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.001})
+        Rota.export_rota_to_html("basic_balancing")
+
+        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)
+            assert worker_shifts["a"] in (22, 23, 24, 25)
+            assert worker_shifts["b"] in (22, 23, 24, 25)
+
+    def test_weighted_shift_balancing(self):
+        Rota = generate_basic_rota(20)
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                balance_weighting=4,
+                workers_required=1,
+                force_as_block=False,
+                constraints=["preclear2", "postclear2"],
+            ),
+            SingleShift(
+                ("group1", "group2"),
+                "b",
+                12.5,
+                days,
+                workers_required=1,
+                force_as_block=False
+            ),
+            SingleShift(
+                ("group1", "group2"),
+                "c",
+                12.5,
+                days[0],
+                workers_required=1,
+                force_as_block=False
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.001})
+        Rota.export_rota_to_html("basic_balancing_weighted_shifts")
+
+        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)
+            assert worker_shifts["a"] in (46,47,48)
+            #assert worker_shifts["b"] in (22, 23, 24, 25)
+
+    def test_weighted_shift_balancing2(self):
+        Rota = generate_basic_rota(23)
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                #balance_weighting=4,
+                workers_required=1,
+                force_as_block=False,
+                #constraints=["preclear2", "postclear2"],
+            ),
+            SingleShift(
+                ("group1", "group2"),
+                "b",
+                12.5,
+                days,
+                balance_weighting=4,
+                workers_required=1,
+                force_as_block=False
+            ),
+            SingleShift(
+                ("group1", "group2"),
+                "c",
+                12.5,
+                days[0],
+                workers_required=1,
+                force_as_block=False
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.001})
+        Rota.export_rota_to_html("basic_balancing_weighted_shifts2")
+
+        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)
+            #assert worker_shifts["a"] in (46,47,48)
+            #assert worker_shifts["c"] in (46,47,48)
+            assert worker_shifts["b"] in (52,53,54,55)
+
+    def test_weighted_shift_balancing3(self):
+        Rota = generate_basic_rota(23)
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                #balance_weighting=4,
+                workers_required=1,
+                force_as_block=False,
+                #constraints=["preclear2", "postclear2"],
+            ),
+            SingleShift(
+                ("group1", "group2"),
+                "b",
+                12.5,
+                days,
+                balance_weighting=4,
+                workers_required=1,
+                force_as_block=False
+            ),
+            SingleShift(
+                ("group1", "group2"),
+                "c",
+                12.5,
+                days[0],
+                balance_weighting=4,
+                workers_required=1,
+                force_as_block=False
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("basic_balancing_weighted_shifts3")
+
+        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)
+            #assert worker_shifts["a"] in (46,47,48)
+            assert worker_shifts["c"] in (7,8)
+            assert worker_shifts["b"] in (52,53,54,55)
+
+    def test_weighted_shift_balancing4(self):
+        Rota = generate_basic_rota(10)
+
+        Rota.add_worker(
+            Worker(Rota, "worker4", "group3", 1),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                balance_offset=99,
+                balance_weighting=8,
+                workers_required=2,
+                force_as_block=False,
+                #constraints=["preclear2", "postclear2"],
+            ),
+            SingleShift(
+                ("group2","group3"),
+                "b",
+                12.5,
+                days,
+                balance_offset=99,
+                #balance_weighting=4,
+                workers_required=1,
+                force_as_block=False
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("basic_balancing_weighted_shifts4")
+
+        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)
+            assert worker_shifts["a"] in (0, 46,47,48)
+
+    def test_weighted_shift_balancing5(self):
+        Rota = generate_basic_rota(10)
+
+        Rota.add_worker(
+            Worker(Rota, "worker4", "group3", 1),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                balance_offset=99,
+                #balance_weighting=8,
+                workers_required=2,
+                force_as_block=False,
+                #constraints=["preclear2", "postclear2"],
+            ),
+            SingleShift(
+                ("group2","group3"),
+                "b",
+                12.5,
+                days,
+                balance_offset=99,
+                balance_weighting=4,
+                workers_required=1,
+                force_as_block=False
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("basic_balancing_weighted_shifts5")
+
+        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)
+            assert worker_shifts["b"] in (0, 35)

test/test_rota.py

@@ -227,7 +227,6 @@ class TestDemoRotaNights:
             "night_weekday",
             12.5,
             days[:4],
-            balance_offset=10,
             workers_required=1,
             constraints=["preclear2", "postclear2"],
             # constraints=["night"],
@@ -241,7 +240,6 @@ class TestDemoRotaNights:
             "night_weekend",
             12.5,
             days[4:],
-            balance_offset=40,
             # balance_weighting=0.5
             constraints=["preclear2", "postclear2"],
             # constraints=["night"],
@@ -252,7 +250,6 @@ class TestDemoRotaNights:
             "twilight",
             12.5,
             days[:5],
-            balance_offset=40,
             # balance_weighting=0.5
             workers_required=2,
             # constraints=["preclear2", "postclear2"],
@@ -268,37 +265,36 @@ class TestDemoRotaNights:
     Rota.build_workers()
     Rota.build_model()
 
-    print(Rota.get_worker_details())
+    solver_options = {"ratio": 0.00, "seconds": 1000, "threads": 10}
-
-    for w in Rota.get_workers():
-        print(w)
-        print(w.get_shift_targets())
-
-    solver_options = {"ratio": 0.1, "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)
-
-    worker_timetable_brief = Rota.get_worker_timetable_brief(
-        show_prefs=False, show_unavailable=False
-    )
-
-    print(worker_timetable_brief)
-    print(Rota.get_shift_summary_dict())
     Rota.export_rota_to_html("test3")
 
-    print(Rota.get_worker_shift_list(worker1))
 
     def test_start_date(self):
         assert self.Rota.start_date == self.start_date
 
 
+    def test_night_assignment(self):
+        shift_summary = self.Rota.get_shift_summary_dict()
+
+        for worker_name in shift_summary:
+            worker_shifts = shift_summary[worker_name]
+
+            worker = self.Rota.get_worker_by_name(worker_name)
+
+            if worker.fte == 40:
+                assert worker_shifts["night_weekday"] in (0,4)
+                assert worker_shifts["night_weekend"] in (0,3)
+            elif worker.fte == 60:
+                assert worker_shifts["night_weekday"] in (0,4,8)
+                assert worker_shifts["night_weekend"] in (0,3,6)
+            else:
+                assert worker_shifts["night_weekday"] in (4,8)
+                assert worker_shifts["night_weekend"] in (3,6)
+
+
 class TestDemoRotaClear:
     weeks_to_rota = 10
     start_date = datetime.date(2022, 3, 7)
@@ -946,10 +942,10 @@ class TestNightUnavailable:
 
     # Add a few workers
     worker1 = Worker(
-        Rota, "worker1", "group1", 1, not_available_to_work=(("15/03/22", "****"),)
+        Rota, "worker1", "group1", 1, not_available_to_work=((datetime.datetime.strptime("15/03/22", "%d/%m/%y").date(), "****"),)
     )
     worker2 = Worker(
-        Rota, "worker2", "group1", 1, not_available_to_work=(("14/03/22", "****"),)
+        Rota, "worker2", "group1", 1, not_available_to_work=((datetime.datetime.strptime("14/03/22", "%d/%m/%y").date(), "****"),)
     )
 
     Rota.add_workers((worker1, worker2))
@@ -1011,10 +1007,10 @@ class TestNightUnavailable:
     def test_assign_split(self):
         self.Rota.shifts = []
         worker3 = Worker(
-            self.Rota, "worker3", "group1", 1, not_available_to_work=(("13/03/22", "****"),)
+            self.Rota, "worker3", "group1", 1, not_available_to_work=((datetime.datetime.strptime("13/03/22", "%d/%m/%y").date(), "****"),)
         )
         worker4 = Worker(
-            self.Rota, "worker4", "group1", 1, not_available_to_work=(("12/03/22", "****"),)
+            self.Rota, "worker4", "group1", 1, not_available_to_work=((datetime.datetime.strptime("12/03/22", "%d/%m/%y").date(), "****"),)
         )
         self.Rota.add_workers((worker3, worker4))
 
@@ -1034,25 +1030,25 @@ class TestNightUnavailable:
         for worker in summary:
             assert summary[worker]["night_weekend"] == 15
 
-    def test_assign_split_night_constraint(self):
+#    def test_assign_split_night_constraint(self):
-        self.Rota.shifts = []
+#        self.Rota.shifts = []
-        worker3 = Worker(
+#        worker3 = Worker(
-            self.Rota, "worker3", "group1", 1, not_available_to_work=(("13/03/22", "****"),)
+#            self.Rota, "worker3", "group1", 1, not_available_to_work=((datetime.datetime.strptime("13/03/22", "%d/%m/%y").date(), "****"),)
-        )
+#        )
-        worker4 = Worker(
+#        worker4 = Worker(
-            self.Rota, "worker4", "group1", 1, not_available_to_work=(("12/03/22", "****"),)
+#            self.Rota, "worker4", "group1", 1, not_available_to_work=((datetime.datetime.strptime("12/03/22", "%d/%m/%y").date(), "****"),)
-        )
+#        )
-        self.Rota.add_workers((worker3, worker4))
+#        self.Rota.add_workers((worker3, worker4))
-
+#
-        self.Rota.add_shifts(
+#        self.Rota.add_shifts(
-            SingleShift(
+#            SingleShift(
-                ("group1",), "night_weekend", 12.5, days[5:], workers_required=3, constraints=["night"]
+#                ("group1",), "night_weekend", 12.5, days[5:], workers_required=3, constraints=["night"]
-            ),
+#            ),
-        )
+#        )
-
+#
-        self.Rota.build_and_solve(options={"ratio": 0.00})
+#        self.Rota.build_and_solve(options={"ratio": 0.00})
-
+#
-        assert self.Rota.results.solver.status == "warning"
+#        assert self.Rota.results.solver.status == "warning"
 
 if __name__ == "__main__":
     t = TestLimitConstraints()

test/test_shifts.py

@@ -0,0 +1,77 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days
+
+from rota.workers import Worker
+
+import itertools
+
+
+def generate_basic_rota(weeks_to_rota=10):
+    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(Rota, "worker1", "group1", 1),
+            Worker(Rota, "worker2", "group1", 1),
+        ]
+    )
+
+    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 TestShifts:
+    def test_duplicate_shifts(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1"),
+                "a",
+                12.5,
+                days,
+            ),
+        )
+
+        with pytest.raises(InvalidShift):
+            Rota.build_and_solve(options={"ratio": 0.000})

test/test_weekend_padding.py

@@ -1,5 +1,3 @@
-from copy import deepcopy
-from black import main
 import pytest
 from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
 

test/test_worker_requests.py

@@ -0,0 +1,499 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days
+
+from rota.workers import Worker
+
+import itertools
+
+
+def generate_basic_rota(weeks_to_rota=10):
+    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(Rota, "worker1", "group1", 1),
+            Worker(Rota, "worker2", "group1", 1),
+        ]
+    )
+
+    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_preferences_not_to_work(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                pref_not_to_work=(list(date_generator(datetime.date(2022, 3, 7), 5))),
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("preferences_not_to_work")
+
+        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)
+
+            assert worker_shifts["a"] in (46, 47)
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            if worker.name == "worker3":
+                assert shift_string.startswith("-----")
+            else:
+                assert shift_string.startswith("aaaaa")
+
+    def test_preferences_not_to_work2(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                pref_not_to_work=(list(date_generator(datetime.date(2022, 3, 7), 55))),
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("preferences_not_to_work2")
+
+        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)
+
+            assert worker_shifts["a"] in (46, 47)
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            if worker.name == "worker3":
+                assert shift_string.endswith("a" * 15)
+
+    def test_preferences_not_to_work_all_shifts(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                pref_not_to_work=(list(date_generator(datetime.date(2022, 3, 7), 70))),
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("preferences_not_to_work_all_shifts")
+
+        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)
+
+            assert worker_shifts["a"] in (46, 47)
+
+            # shift_string = Rota.get_worker_shift_list_string(worker)
+
+            # if worker.name == "worker3":
+            #    assert shift_string.startswith("-----")
+            # else:
+            #    assert shift_string.startswith("aaaaa")
+
+    def test_unavailable_to_work(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                not_available_to_work=(
+                    [
+                        (i, "NOT AROUND")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 23))
+                    ]
+                ),
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("unavailable_to_work")
+
+        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)
+
+            assert worker_shifts["a"] in (46, 47)
+
+            # shift_string = Rota.get_worker_shift_list_string(worker)
+
+            # if worker.name == "worker3":
+            #    assert shift_string.startswith("-----")
+            # else:
+            #    assert shift_string.startswith("aaaaa")
+
+    def test_unavailable_to_work2(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                not_available_to_work=(
+                    [
+                        (i, "NOT AROUND")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 25))
+                    ]
+                ),
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=2,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("unavailable_to_work2")
+
+        assert Rota.results.solver.status == "warning"
+        assert Rota.results.solver.termination_condition == "infeasible"
+
+    def test_unavailable_to_work3(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+
+        Rota.constraint_options["balance_bank_holidays"] = False
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                not_available_to_work=(
+                    [
+                        (i, "NOT AROUND")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 23))
+                    ]
+                ),
+                pref_not_to_work=[datetime.date(2022, 5, 15)],
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=2,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("unavailable_to_work3")
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+        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)
+
+            assert worker_shifts["a"] in (46, 47)
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            if worker.name == "worker3":
+               assert shift_string.startswith("-"*23)
+               assert shift_string.endswith("-")
+
+    def test_unavailable_to_work2(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+
+        Rota.add_workers(
+            [Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                not_available_to_work=(
+                    [
+                        (i, "NOT AROUND")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 25))
+                    ]
+                ),
+            ),
+            Worker(
+                Rota,
+                "worker4",
+                "group1",
+                1,
+                not_available_to_work=(
+                    [
+                        (i, "NOT AROUND")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 25))
+                    ]
+                ),
+            ),]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=3,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("unavailable_to_work4")
+
+        assert Rota.results.solver.status == "warning"
+        assert Rota.results.solver.termination_condition == "infeasible"
+
+    def test_work_requests(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+        Rota.constraint_options["balance_bank_holidays"] = False
+
+        Rota.add_workers(
+            [Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                work_requests=(
+                    [
+                        (i, "a")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
+                    ]
+                ),
+            ),
+            Worker(
+                Rota,
+                "worker4",
+                "group1",
+                1,
+                work_requests=(
+                    [
+                        (i, "a")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
+                    ]
+                ),
+            ),]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=2,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("work_requests")
+
+        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)
+
+            assert worker_shifts["a"] == 35
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            if worker.name in ("worker3", "worker4"):
+               assert shift_string.startswith("a"*35)
+
+    def test_work_requests_invalid(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+        Rota.constraint_options["balance_bank_holidays"] = False
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker3",
+                "group1",
+                1,
+                work_requests=(
+                    [
+                        (i, "b")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
+                    ]
+                ),
+            ),
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "a",
+                12.5,
+                days,
+                workers_required=1,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+
+        with pytest.raises(InvalidShift):
+            Rota.build_and_solve(options={"ratio": 0.000})
+        
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "c",
+                12.5,
+                days,
+                workers_required=1,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        with pytest.raises(InvalidShift):
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+        Rota.add_shifts(
+            SingleShift(
+                ("group1", "group2"),
+                "b",
+                12.5,
+                days,
+                workers_required=1,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+        Rota.add_worker(
+            Worker(
+                Rota,
+                "worker4",
+                "group1",
+                1,
+                work_requests=(
+                    [
+                        (i, "d")
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
+                    ]
+                ),
+            ),
+        )
+
+        with pytest.raises(InvalidShift):
+            Rota.build_and_solve(options={"ratio": 0.000})