Fix file imports - fix tests with redirs

1 files changed, 20 insertions, 3 deletions

diff --git a/ishtar_common/tests.py b/ishtar_common/tests.py
index 765d74312..8708d90bd 100644
--- a/ishtar_common/tests.py
+++ b/ishtar_common/tests.py
@@ -414,6 +414,9 @@ class WizardTest(object):
     condition_dict = None
     form_datas = []
     test_back = True
+    redirect_url = None
+    model = None
+    current_id_key = "pk"
 
     def setUp(self):
         self.username, self.password, self.user = create_superuser()
@@ -502,9 +505,23 @@ class WizardTest(object):
             )
         if idx == len(self.steps) - 1:
             #  last form
-            self.assertRedirects(
-                response,
-                '/{}/done'.format(self.url_uri))
+            if not self.redirect_url:
+                redirect_url = '/{}/done'.format(self.url_uri)
+            else:
+                dct = {
+                    "url_name": self.url_name,
+                    "url_uri": self.url_uri
+                }
+                form_key = 'selec-' + self.url_name
+                if form_key in form_data and self.current_id_key in form_data[
+                        form_key]:
+                    dct["current_id"] = form_data[form_key][self.current_id_key]
+                if self.model:
+                    q = self.model.objects
+                    if q.count():
+                        dct["last_id"] = q.order_by("-pk").all()[0].pk
+                redirect_url = self.redirect_url.format(**dct)
+            self.assertRedirects(response, redirect_url)
         return response
 
     def test_wizard(self):