fix multishifts

test/test_multiple_shifts_per_day.py

@@ -31,58 +31,56 @@ def test_single_shift_per_day_default(basic_rota):
 def test_allow_two_shifts_per_day(basic_rota):
     Rota = basic_rota
     Rota.add_shifts(
-        SingleShift(sites=["site1"], name="a", length=8, days=days[:1], workers_required=1),
-        SingleShift(sites=["site1"], name="b", length=8, days=days[:1], workers_required=1),
+        SingleShift(sites=["site1"], name="a", length=8, days=days[:2], workers_required=1),
+        SingleShift(sites=["site1"], name="b", length=8, days=days[:2], workers_required=1),
     )
     Rota.allow_shifts_together("a", "b")
     Rota.build_and_solve()
     Rota.export_rota_to_html("test_allow_two_shifts", folder="tests")
     assert Rota.results.solver.status == "ok"
 
-    ## Now both shifts can be assigned together
-    #for week, day in Rota.get_week_day_combinations():
-    #    assigned = [
-    #        shift for shift in Rota.get_shift_names_by_week_day(week, day)
-    #        if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5
-    #    ]
-    #    assert len(assigned) <= 2
 
-def test_only_specified_shifts_can_be_double_assigned(basic_rota):
-    Rota = basic_rota
-    Rota.add_shifts(
-        SingleShift(sites=["site1"], name="a", length=8, days=days[:5], workers_required=1),
-        SingleShift(sites=["site1"], name="b", length=8, days=days[:5], workers_required=1),
-        SingleShift(sites=["site1"], name="c", length=8, days=days[:5], workers_required=1),
-    )
-    Rota.allow_shifts_together("a", "b")
-    Rota.build_and_solve()
-    # Only "a" and "b" can be assigned together, never "a"+"c" or "b"+"c"
+    # Now both shifts can be assigned together
     for week, day in Rota.get_week_day_combinations():
         assigned = [
             shift for shift in Rota.get_shift_names_by_week_day(week, day)
             if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5
         ]
-        if set(["a", "b"]).issubset(set(Rota.get_shift_names_by_week_day(week, day))):
-            assert len(assigned) <= 2
-        else:
-            assert len(assigned) <= 1
+        assert len(assigned) <= 2
+
+def test_only_specified_shifts_can_be_double_assigned(basic_rota):
+    Rota = basic_rota
+    worker2 = Worker(name="worker2", site="site1", grade=1)
+    Rota.add_worker(worker2)
+    #Rota.add_worker(worker3)
+    Rota.add_shifts(
+        SingleShift(sites=["site1"], name="a", length=8, days=days[:3], workers_required=1),
+        SingleShift(sites=["site1"], name="b", length=8, days=days[:3], workers_required=1),
+        SingleShift(sites=["site1"], name="c", length=8, days=days[:3], workers_required=1),
+    )
+    Rota.allow_shifts_together("a", "b")
+    Rota.build_and_solve()
+    Rota.export_rota_to_html("test_only_specified_shifts_can_be_double_assigned", folder="tests")
+    assert Rota.results.solver.status == "ok"
+    # Only "a" and "b" can be assigned together, never "a"+"c" or "b"+"c"
+    for worker in Rota.workers:
+        for week, day in Rota.get_week_day_combinations():
+            assigned = set([
+                shift for shift in Rota.get_shift_names_by_week_day(week, day)
+                if Rota.model.works[worker.id, week, day, shift].value > 0.8
+            ])
+
+            assert assigned != set(["a", "c"]) and assigned != set(["b", "c"])
 
 def test_triple_shifts_allowed(basic_rota):
     Rota = basic_rota
     Rota.add_shifts(
-        SingleShift(sites=["site1"], name="a", length=8, days=days[:5], workers_required=1),
-        SingleShift(sites=["site1"], name="b", length=8, days=days[:5], workers_required=1),
-        SingleShift(sites=["site1"], name="c", length=8, days=days[:5], workers_required=1),
+        SingleShift(sites=["site1"], name="a", length=8, days=days[2:4], workers_required=1),
+        SingleShift(sites=["site1"], name="b", length=8, days=days[2:4], workers_required=1),
+        SingleShift(sites=["site1"], name="c", length=8, days=days[2:4], workers_required=1),
     )
     Rota.allow_shifts_together("a", "b", "c")
     Rota.build_and_solve()
-    # All three can be assigned together if allowed
-    for week, day in Rota.get_week_day_combinations():
-        assigned = [
-            shift for shift in Rota.get_shift_names_by_week_day(week, day)
-            if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5
-        ]
-        if set(["a", "b", "c"]).issubset(set(Rota.get_shift_names_by_week_day(week, day))):
-            assert len(assigned) <= 3
-        else:
-            assert len(assigned) <= 1
+    Rota.export_rota_to_html("test_triple_shifts_allowed", folder="tests")
+    assert Rota.results.solver.status == "ok"
+