From e34d9f446b4f63db11ad58d9871a8b8c68061393 Mon Sep 17 00:00:00 2001
From: Parbhat Puri <parbhatpuri17@gmail.com>
Date: Wed, 7 Aug 2019 09:02:28 +0000
Subject: [PATCH] GH-717: Update test suite for new field stage in application,
 lab forms

---
 opentech/apply/funds/tests/factories/models.py |  9 ++++++---
 opentech/apply/funds/tests/test_admin_form.py  |  8 +++++---
 opentech/apply/funds/tests/test_admin_views.py | 14 ++++++++++----
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/opentech/apply/funds/tests/factories/models.py b/opentech/apply/funds/tests/factories/models.py
index d31bd9bf20..6b348db8f0 100644
--- a/opentech/apply/funds/tests/factories/models.py
+++ b/opentech/apply/funds/tests/factories/models.py
@@ -90,10 +90,11 @@ def parent(self, create, extracted_parent, **parent_kwargs):
     @factory.post_generation
     def forms(self, create, extracted, **kwargs):
         if create:
-            for _ in self.workflow.stages:
+            for index, _ in enumerate(self.workflow.stages, 1):
                 # Generate a form based on all defined fields on the model
                 ApplicationBaseFormFactory(
                     application=self,
+                    stage=index,
                     **kwargs,
                 )
                 ApplicationBaseReviewForm(
@@ -160,10 +161,11 @@ def parent(self, create, extracted_parent, **parent_kwargs):
     @factory.post_generation
     def forms(self, create, extracted, **kwargs):
         if create:
-            for _ in self.workflow.stages:
+            for index, _ in enumerate(self.workflow.stages, 1):
                 # Generate a form based on all defined fields on the model
                 RoundBaseFormFactory(
                     round=self,
+                    stage=index,
                     **kwargs,
                 )
                 RoundBaseReviewFormFactory(
@@ -197,10 +199,11 @@ class Meta:
     @factory.post_generation
     def forms(self, create, extracted, **kwargs):
         if create:
-            for _ in self.workflow.stages:
+            for index, _ in enumerate(self.workflow.stages, 1):
                 # Generate a form based on all defined fields on the model
                 LabBaseFormFactory(
                     lab=self,
+                    stage=index,
                     **kwargs,
                 )
                 LabBaseReviewFormFactory(
diff --git a/opentech/apply/funds/tests/test_admin_form.py b/opentech/apply/funds/tests/test_admin_form.py
index aa0c78f0ef..244ece50b1 100644
--- a/opentech/apply/funds/tests/test_admin_form.py
+++ b/opentech/apply/funds/tests/test_admin_form.py
@@ -38,9 +38,11 @@ def formset_base(field, total, delete, factory, same=False, form_stage_info=None
     return base_data
 
 
-def form_data(number_application_forms=0, number_review_forms=0, delete=0, stages=1, same_forms=False, form_stage_info=[1]):
-    form_data = formset_base('forms', number_application_forms, delete, same=same_forms, factory=ApplicationFormFactory, form_stage_info=form_stage_info)
-    review_form_data = formset_base('review_forms', number_review_forms, False, same=same_forms, factory=ReviewFormFactory)
+def form_data(num_appl_forms=0, num_review_forms=0, delete=0, stages=1, same_forms=False, form_stage_info=[1]):
+    form_data = formset_base(
+        'forms', num_appl_forms, delete, same=same_forms, factory=ApplicationFormFactory,
+        form_stage_info=form_stage_info)
+    review_form_data = formset_base('review_forms', num_review_forms, False, same=same_forms, factory=ReviewFormFactory)
     form_data.update(review_form_data)
 
     fund_data = factory.build(dict, FACTORY_CLASS=FundTypeFactory)
diff --git a/opentech/apply/funds/tests/test_admin_views.py b/opentech/apply/funds/tests/test_admin_views.py
index 9bfc4cc438..4cda56b836 100644
--- a/opentech/apply/funds/tests/test_admin_views.py
+++ b/opentech/apply/funds/tests/test_admin_views.py
@@ -17,11 +17,12 @@ def setUpTestData(cls):
         cls.user = SuperUserFactory()
         cls.home = ApplyHomePageFactory()
 
-    def create_page(self, forms=1, same_forms=False):
+    def create_page(self, appl_forms=1, review_forms=1, stages=1, same_forms=False, form_stage_info=[1]):
         self.client.force_login(self.user)
         url = reverse('wagtailadmin_pages:add', args=('funds', 'fundtype', self.home.id))
 
-        data = form_data(forms, same_forms=same_forms, stages=forms)
+        data = form_data(
+            appl_forms, review_forms, same_forms=same_forms, stages=stages, form_stage_info=form_stage_info)
         data['action-publish'] = True
 
         response = self.client.post(url, data=data, secure=True, follow=True)
@@ -38,12 +39,17 @@ def test_can_create_fund(self):
         self.assertEqual(fund.review_forms.count(), 1)
 
     def test_can_create_multi_phase_fund(self):
-        fund = self.create_page(2)
+        fund = self.create_page(2, 2, stages=2, form_stage_info=[1, 2])
         self.assertEqual(fund.forms.count(), 2)
         self.assertEqual(fund.review_forms.count(), 2)
 
+    def test_can_create_multiple_forms_second_stage_in_fund(self):
+        fund = self.create_page(4, 2, stages=2, form_stage_info=[1, 2, 2, 2])
+        self.assertEqual(fund.forms.count(), 4)
+        self.assertEqual(fund.review_forms.count(), 2)
+
     def test_can_create_multi_phase_fund_reuse_forms(self):
-        fund = self.create_page(2, same_forms=True)
+        fund = self.create_page(2, 2, same_forms=True, stages=2, form_stage_info=[1, 2])
         self.assertEqual(fund.forms.count(), 2)
         self.assertEqual(fund.review_forms.count(), 2)