Merge pull request #129 from 2gis/create-step-for-response-even-if-ex…

…ception-happens

Create step for response even if exception happens

core/sessions.py

@@ -48,7 +48,7 @@ def __init__(self, method, url="/", headers=None, data=""):
 
     def __repr__(self):
         return "<RequestHelper method:%s url:%s headers:%s body:%s>" % (
-            self.method, self.url, self.headers, self.body)
+            self.method, self.url, self.headers, self.data)
 
 
 def update_log_step(log_step, message=None, control_line=None):

tests/unit/test_server.py

@@ -643,3 +643,86 @@ def test_max_count_with_run_new_request_during_prevm_is_not_ready(self):
 
         self.assertEqual(200, response.status)
         self.assertEqual(1, self.pool.count())
+
+
+class TestSessionSteps(BaseTestServer):
+    def setUp(self):
+        setup_config('data/config.py')
+
+        self.address = ("localhost", 9001)
+
+        super(TestSessionSteps, self).setUp()
+
+        self.desired_caps = {
+            'desiredCapabilities': {
+                'platform': self.vmmaster.app.platforms.platforms.keys()[0]
+            }
+        }
+
+        self.ctx = self.vmmaster.app.app_context()
+        self.ctx.push()
+
+    def tearDown(self):
+        del self.vmmaster
+        server_is_down(self.address)
+        self.ctx.pop()
+
+    def test_add_first_two_steps(self):
+        """
+        - exception while waiting endpoint
+        Expected: session was created, session_step was created
+        """
+
+        def raise_exception(dc):
+            raise Exception('something ugly happened in new_vm')
+
+        with patch(
+            'vmpool.endpoint.new_vm', Mock(side_effect=raise_exception)
+        ), patch(
+            'core.sessions.Session.add_session_step', Mock()
+        ) as add_step_mock:
+            response = new_session_request(self.address, self.desired_caps)
+
+        self.assertEqual(500, response.status)
+        self.assertIn('something ugly happened in new_vm', response.content)
+
+        self.assertEqual(add_step_mock.call_count, 2)
+
+    @patch(
+        'vmmaster.webdriver.commands.ping_vm',
+        new=Mock(side_effect=ping_vm_true_mock)
+    )
+    @patch(
+        'vmmaster.webdriver.commands.selenium_status',
+        new=Mock(
+            __name__="selenium_status",
+            return_value=(200, {}, json.dumps({'status': 0}))
+        )
+    )
+    def test_always_create_response_for_sub_step(self):
+        """
+        - exception while executing make_request
+        Expected: session failed, sub_steps was created
+        """
+
+        def raise_exception(*args, **kwargs):
+            raise Exception('something ugly happened in make_request')
+
+        def new_vm_mock(arg):
+            yield Mock(ip=1)
+
+        with patch(
+            'vmpool.endpoint.new_vm', Mock(side_effect=new_vm_mock)
+        ), patch(
+            'core.sessions.Session.make_request',
+            Mock(__name__="make_request", side_effect=raise_exception)
+        ), patch(
+            'core.sessions.Session.add_sub_step', Mock()
+        ) as add_sub_step_mock:
+            response = new_session_request(self.address, self.desired_caps)
+
+        self.assertEqual(500, response.status)
+        self.assertIn(
+            'something ugly happened in make_request', response.content)
+
+        self.assertEqual(add_sub_step_mock.call_count, 2)

vmmaster/webdriver/__init__.py

@@ -96,14 +96,6 @@ def delete_session(session_id):
 def create_session():
     if current_app.running:
         session = helpers.get_session()
-        control_line = "%s %s %s" % (
-            request.method, request.path,
-            request.headers.environ['SERVER_PROTOCOL']
-        )
-        session.add_session_step(
-            control_line=control_line,
-            body=str(request.data)
-        )
         commands.replace_platform_with_any(request)
         status, headers, body = commands.start_session(request, session)
         return helpers.form_response(status, headers, body)

vmmaster/webdriver/commands.py

@@ -6,6 +6,7 @@
 import websocket
 import thread
 
+from traceback import format_exc
 from core import utils
 from core.utils import network_utils
 from core.utils import generator_wait_for
@@ -25,16 +26,23 @@ def wrapper(port, request, *args, **kwargs):
             control_line="%s %s" % (request.method, request.url),
             body=request.data)
 
-        for status, headers, body in func(port, request, *args, **kwargs):
-            yield status, headers, body
-
-        content_to_log = utils.remove_base64_screenshot(body)
+        try:
+            for status, headers, body in func(port, request, *args, **kwargs):
+                yield status, headers, body
+        except Exception as e:
+            session.add_sub_step(
+                control_line='500',
+                body=format_exc()
+            )
+            raise e
+        else:
+            content_to_log = utils.remove_base64_screenshot(body)
 
-        session.add_sub_step(
-            control_line=str(status),
-            body=content_to_log)
+            session.add_sub_step(
+                control_line=str(status),
+                body=content_to_log)
 
-        yield status, headers, body
+            yield status, headers, body
 
     return wrapper
 

vmmaster/webdriver/helpers.py

@@ -146,16 +146,15 @@ def transparent():
 
 
 def vmmaster_agent(command):
-    session = current_app.sessions.get_session(request.session.id)
+    session = request.session
     swap_session(request, session.selenium_session)
     code, headers, body = command(request, session)
     swap_session(request, session.selenium_session)
     return code, headers, body
 
 
 def internal_exec(command):
-    session = current_app.sessions.get_session(request.session.id)
-    code, headers, body = command(request, session)
+    code, headers, body = command(request, request.session)
     return code, headers, body
 
 
@@ -182,7 +181,14 @@ def get_session():
     request.session = session
     log.info("New session %s (%s) for %s" %
              (str(session.id), session.name, str(dc)))
-
+    control_line = "%s %s %s" % (
+        request.method, request.path,
+        request.headers.environ['SERVER_PROTOCOL']
+    )
+    session.add_session_step(
+        control_line=control_line,
+        body=str(request.data)
+    )
     yield session
     for vm in endpoint.new_vm(dc):
         session.endpoint = vm