Refactor: "Visit page for poll" step

1 files changed, 12 insertions, 4 deletions

diff --git a/webapp/webapp/features/steps.py b/webapp/webapp/features/steps.py
index fb398d4..5f57673 100644
--- a/webapp/webapp/features/steps.py
+++ b/webapp/webapp/features/steps.py
@@ -116,11 +116,19 @@ def i_should_see_a_message_that_says_text(step, text):
     assert_true(b.is_text_present(text))
 
 
+def get_poll_by_name(poll_name):
+    return filter(lambda p: p.name == poll_name, Poll.all())[0]
+
+
 @step(u'When I visit the "([^"]*)" page for "([^"]*)" poll')
 def visit_page_for_poll(step, page_name, poll_name):
-    page_list = {'Resultados': 'results'}
-    page = page_list[page_name]
-    poll_id = world.poll_id.__str__()
+    poll_id = getattr(world, "poll_id", None)
+    poll_id = str(poll_id) if poll_id else str(get_poll_by_name(poll_name).id)
+    page_list = {
+        'Resultados': [poll_id, 'results'],
+        'Modificar estructura': ['structure', poll_id]
+    }
+    page_args = page_list[page_name]
     b = world.browser
-    url = '/manager/polls/{poll_id}/{page}/'.format(poll_id=poll_id, page=page)
+    url = '/manager/polls/' + "/".join(page_args)
     b.visit(django_url(url))