Correctly handle the last stage in the Workflow

opentech/apply/tests.py

@@ -28,6 +28,10 @@ class TestWorkflowCreation(SimpleTestCase):
         workflow = WorkflowFactory(num_stages=1)
         self.assertEqual(workflow.next(), workflow.stages[0])
 
+    def test_returns_none_if_no_next(self):
+        workflow = WorkflowFactory(num_stages=1)
+        self.assertEqual(workflow.next(workflow.stages[0]), None)
+
 
 class TestStageCreation(SimpleTestCase):
     def test_can_create_stage(self):

opentech/apply/workflow.py

-from typing import Iterator, Iterable
+from typing import Iterator, Iterable, Union
 
 from django.forms import Form
 
@@ -13,8 +13,11 @@ class Workflow(Iterable['Stage']):
     def __iter__(self) -> Iterator['Stage']:
         yield from self.stages
 
-    def next(self):
-        return self.stages[0]
+    def next(self, current_stage: Union['Stage', None]=None) -> Union['Stage', None]:
+        if not current_stage:
+            return self.stages[0]
+
+        return None
 
 
 class Stage: