add worker requirements

test/test_workers.py

@@ -2,23 +2,29 @@ import pytest
 from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, days
 
 import datetime
-from rota.workers import Worker
+from rota.workers import NotAvailableToWork, Worker, generate_not_available_to_works
 
+def generate_basic_rota(weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)):
 
-def setup_basic_rota():
-    weeks_to_rota = 10
-    start_date = datetime.date(2022, 3, 7)
-
-    return RotaBuilder(
+    Rota = RotaBuilder(
         start_date,
         weeks_to_rota=weeks_to_rota,
     )
 
+    # Add a few workers
+    if workers:
+        Rota.add_workers(
+            [
+                Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)
+            ]
+        )
+
+    return Rota
 
 class TestWorkers:
     def test_no_workers(self):
         """Check error raised if attempt to solve with no workers"""
-        Rota = setup_basic_rota()
+        Rota = generate_basic_rota()
 
         Rota.add_shifts(
             SingleShift(
@@ -43,7 +49,7 @@ class TestWorkers:
 
     def test_workers_duplicate_name(self):
 
-        Rota = setup_basic_rota()
+        Rota = generate_basic_rota()
 
         workers = []
         workers.append(Worker(name="worker1", site="group1", grade=1, fte=100))
@@ -325,7 +331,7 @@ class TestWorkers:
         Rota.export_rota_to_html("worker_oops")
 
     def test_worker_pairs(self):
-        Rota = setup_basic_rota()
+        Rota = generate_basic_rota()
 
         Rota.add_shifts(
             SingleShift(
@@ -395,4 +401,42 @@ class TestWorkers:
         Rota.export_rota_to_html("worker_pairs")
 
         assert Rota.results.solver.status == "ok"
-        assert Rota.get_workers_total_shifts()["worker3"] == 50
+        assert Rota.get_workers_total_shifts()["worker3"] == 50
+
+    def test_not_available_to_work(self):
+        
+        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("-")