From f8ba6836b43364925ac1307b80ef46b9dbdaa04f Mon Sep 17 00:00:00 2001
From: Stanislav Zubov <s.zubov@2Gis.ru>
Date: Sat, 17 Oct 2015 13:39:15 +0600
Subject: [PATCH] Separate wd routes for /vmmaster/<command>

---
 tests/unit/helpers.py          |  4 +--
 tox.ini                        |  2 +-
 vmmaster/webdriver/__init__.py | 55 ++++++++++++++++++++++------------
 3 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/tests/unit/helpers.py b/tests/unit/helpers.py
index 5459e593..cf8da13b 100644
--- a/tests/unit/helpers.py
+++ b/tests/unit/helpers.py
@@ -77,13 +77,13 @@ def get_session_request(address, session):
 
 def run_script(address, session, script):
     return request("%s:%s" % address, "POST",
-                   "/wd/hub/session/%s/runScript" % str(session),
+                   "/wd/hub/session/%s/vmmaster/runScript" % str(session),
                    body=json.dumps({"script": script}))
 
 
 def vmmaster_label(address, session, label=None):
     return request("%s:%s" % address, "POST",
-                   "/wd/hub/session/%s/vmmasterLabel" % str(session),
+                   "/wd/hub/session/%s/vmmaster/vmmasterLabel" % str(session),
                    body=json.dumps({"label": label}))
 
 from core.utils.network_utils import get_socket
diff --git a/tox.ini b/tox.ini
index 3bb62ae1..7b3cf670 100644
--- a/tox.ini
+++ b/tox.ini
@@ -15,7 +15,7 @@ deps =  selenium==2.46.0
         mock==1.0.1
         flake8==2.4.1
         coverage==3.7.1
-        git+https://github.com/2gis/vmmaster-client
+        git+https://github.com/2gis/vmmaster-client@urls-for-custom-commands
         git+https://github.com/2gis/lode_runner
         {[testenv:base]deps}
 
diff --git a/vmmaster/webdriver/__init__.py b/vmmaster/webdriver/__init__.py
index 7b9b7bb7..16ca2549 100644
--- a/vmmaster/webdriver/__init__.py
+++ b/vmmaster/webdriver/__init__.py
@@ -111,25 +111,7 @@ def get_session(session_id):
     return helpers.form_response(status, headers, body)
 
 
-@webdriver.route("/session/<string:session_id>/<path:url>",
-                 methods=['GET', 'POST', 'DELETE'])
-def proxy_request(session_id, url=None):
-    request.session = current_app.sessions.get_session(session_id)
-
-    if url:
-        last = url.split("/")[-1]
-    else:
-        last = None
-
-    if last in commands.AgentCommands:
-        status, headers, body = helpers.vmmaster_agent(
-            commands.AgentCommands[last])
-    elif last in commands.InternalCommands:
-        status, headers, body = helpers.internal_exec(
-            commands.InternalCommands[last])
-    else:
-        status, headers, body = helpers.transparent()
-
+def take_screenshot(status, body):
     words = ["url", "click", "execute", "keys", "value"]
     only_screenshots = ["element", "execute_async"]
     parts = request.path.split("/")
@@ -140,4 +122,39 @@ def proxy_request(session_id, url=None):
         utils.to_thread(
             helpers.take_screenshot_from_response(request.session, body))
 
+
+@webdriver.route(
+    "/session/<string:session_id>/vmmaster/runScript", methods=['POST']
+)
+def agent_command(session_id):
+    request.session = current_app.sessions.get_session(session_id)
+
+    status, headers, body = helpers.vmmaster_agent(
+        commands.AgentCommands['runScript'])
+
+    take_screenshot(status, body)
+    return helpers.form_response(status, headers, body)
+
+
+@webdriver.route(
+    "/session/<string:session_id>/vmmaster/vmmasterLabel", methods=['POST']
+)
+def vmmaster_command(session_id):
+    request.session = current_app.sessions.get_session(session_id)
+
+    status, headers, body = helpers.internal_exec(
+        commands.InternalCommands['vmmasterLabel'])
+
+    take_screenshot(status, body)
+    return helpers.form_response(status, headers, body)
+
+
+@webdriver.route("/session/<string:session_id>/<path:url>",
+                 methods=['GET', 'POST', 'DELETE'])
+def proxy_request(session_id, url=None):
+    request.session = current_app.sessions.get_session(session_id)
+
+    status, headers, body = helpers.transparent()
+
+    take_screenshot(status, body)
     return helpers.form_response(status, headers, body)