diff --git a/opentech/apply/tests/test_workflow.py b/opentech/apply/tests/test_workflow.py index 72f9ad0fcf2937dca727ee51c4f193406f527231..213420444ec4cff22b87e566ca1ea46f38abb79c 100644 --- a/opentech/apply/tests/test_workflow.py +++ b/opentech/apply/tests/test_workflow.py @@ -43,6 +43,11 @@ class TestStageCreation(SimpleTestCase): self.assertEqual(stage.form, form) self.assertEqual(stage.phases, phases) + def test_can_iterate_through_phases(self): + stage = StageFactory() + for phase, check in zip(stage, stage.phases): + self.assertEqual(phase, check) + class TestPhaseCreation(SimpleTestCase): def test_can_create_phase(self): diff --git a/opentech/apply/workflow.py b/opentech/apply/workflow.py index 882208b35b5e4d5a275ccc59e927808b33688e94..414978eb59b430dcc82ce704020180e2acaa8311 100644 --- a/opentech/apply/workflow.py +++ b/opentech/apply/workflow.py @@ -31,6 +31,9 @@ class Stage: self.form = form self.phases = phases + def __iter__(self) -> Iterator['Phase']: + yield from self.phases + class Phase: def __init__(self, name: str) -> None: