[chromedriver] Bidi WebSocket connection(III)

Differentiate passing host for new session and
pass host instead of session id so that we can
use it to construct WebSocketUrl potentially

Bug: chromedriver:3588
Change-Id: Ib7a0dda3daf816b36ecf6220140dc906e2040487
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2391903Reviewed-by: John Chen <johnchen@chromium.org>
Commit-Queue: Shengfa Lin <shengfa@google.com>
Cr-Commit-Position: refs/heads/master@{#806369}

chrome/test/chromedriver/client/chromedriver.py

@@ -151,7 +151,7 @@ class ChromeDriver(object):
       send_w3c_capability=True, send_w3c_request=True,
       page_load_strategy=None, unexpected_alert_behaviour=None,
       devtools_events_to_log=None, accept_insecure_certs=None,
-      timeouts=None, test_name=None):
+      timeouts=None, test_name=None, web_socket_url=None):
     self._executor = command_executor.CommandExecutor(server_url)
     self._server_url = server_url
     self.w3c_compliant = False
@@ -266,6 +266,9 @@ class ChromeDriver(object):
     if test_name is not None:
       params['goog:testName'] = test_name
 
+    if web_socket_url is not None:
+      params['webSocketUrl'] = web_socket_url
+
     if send_w3c_request:
       params = {'capabilities': {'alwaysMatch': params}}
     else:
@@ -722,3 +725,9 @@ class ChromeDriver(object):
     params = {'authenticatorId': authenticatorId,
               'isUserVerified': isUserVerified}
     return self.ExecuteCommand(Command.SET_USER_VERIFIED, params)
+
+  def GetSessionId(self):
+    if not hasattr(self, '_session_id'):
+      return None
+    return self._session_id
+

chrome/test/chromedriver/commands.cc

@@ -61,16 +61,13 @@ void ExecuteGetStatus(
                std::string(), kW3CDefault);
 }
 
-void ExecuteCreateSession(
-    SessionThreadMap* session_thread_map,
-    const Command& init_session_cmd,
-    const base::DictionaryValue& params,
-    const std::string& session_id,
-    const CommandCallback& callback) {
-  std::string new_id = session_id;
-  if (new_id.empty())
-    new_id = GenerateId();
-  std::unique_ptr<Session> session = std::make_unique<Session>(new_id);
+void ExecuteCreateSession(SessionThreadMap* session_thread_map,
+                          const Command& init_session_cmd,
+                          const base::DictionaryValue& params,
+                          const std::string& host,
+                          const CommandCallback& callback) {
+  std::string new_id = GenerateId();
+  std::unique_ptr<Session> session = std::make_unique<Session>(new_id, host);
   std::unique_ptr<SessionThreadInfo> threadInfo =
       std::make_unique<SessionThreadInfo>(new_id, GetW3CSetting(params));
   if (!threadInfo->thread()->Start()) {

chrome/test/chromedriver/commands.h

@@ -28,12 +28,11 @@ void ExecuteGetStatus(
     const CommandCallback& callback);
 
 // Creates a new session.
-void ExecuteCreateSession(
-    SessionThreadMap* session_thread_map,
-    const Command& init_session_cmd,
-    const base::DictionaryValue& params,
-    const std::string& session_id,
-    const CommandCallback& callback);
+void ExecuteCreateSession(SessionThreadMap* session_thread_map,
+                          const Command& init_session_cmd,
+                          const base::DictionaryValue& params,
+                          const std::string& host,
+                          const CommandCallback& callback);
 
 // Gets all sessions
 void ExecuteGetSessions(

chrome/test/chromedriver/server/http_handler.cc

@@ -1061,8 +1061,11 @@ void HttpHandler::HandleCommand(
                     nullptr, session_id, true);
     return;
   }
-
-  iter->command.Run(params, session_id,
+  // Pass host instead for potential WebSocketUrl if it's a new session
+  iter->command.Run(params,
+                    internal::IsNewSession(*iter)
+                        ? request.GetHeaderValue("host")
+                        : session_id,
                     base::BindRepeating(&HttpHandler::PrepareResponse,
                                         weak_ptr_factory_.GetWeakPtr(),
                                         trimmed_path, send_response_func));
@@ -1353,4 +1356,9 @@ bool MatchesCommand(const std::string& method,
   return true;
 }
 
+bool IsNewSession(const CommandMapping& command) {
+  return command.method == kPost &&
+         command.path_pattern == kNewSessionPathPattern;
+}
+
 }  // namespace internal

chrome/test/chromedriver/server/http_handler.h

@@ -162,6 +162,8 @@ bool MatchesCommand(const std::string& method,
                     std::string* session_id,
                     base::DictionaryValue* out_params);
 
+bool IsNewSession(const CommandMapping& command);
+
 }  // namespace internal
 
 #endif  // CHROME_TEST_CHROMEDRIVER_SERVER_HTTP_HANDLER_H_

chrome/test/chromedriver/session.cc

@@ -75,20 +75,13 @@ Session::Session(const std::string& id)
       mouse_click_timestamp(base::TimeTicks::Now()) {}
 
 Session::Session(const std::string& id, std::unique_ptr<Chrome> chrome)
-    : id(id),
-      w3c_compliant(kW3CDefault),
-      quit(false),
-      detach(false),
-      chrome(std::move(chrome)),
-      sticky_modifiers(0),
-      mouse_position(0, 0),
-      pressed_mouse_button(kNoneMouseButton),
-      implicit_wait(kDefaultImplicitWaitTimeout),
-      page_load_timeout(kDefaultPageLoadTimeout),
-      script_timeout(kDefaultScriptTimeout),
-      strict_file_interactability(false),
-      click_count(0),
-      mouse_click_timestamp(base::TimeTicks::Now()) {}
+    : Session(id) {
+  this->chrome = std::move(chrome);
+}
+
+Session::Session(const std::string& id, const std::string& host) : Session(id) {
+  this->host = host;
+}
 
 Session::~Session() {}
 

chrome/test/chromedriver/session.h

@@ -70,6 +70,7 @@ struct Session {
 
   explicit Session(const std::string& id);
   Session(const std::string& id, std::unique_ptr<Chrome> chrome);
+  Session(const std::string& id, const std::string& host);
   ~Session();
 
   Status GetTargetWindow(WebView** web_view);
@@ -129,6 +130,7 @@ struct Session {
   std::string unhandled_prompt_behavior;
   int click_count;
   base::TimeTicks mouse_click_timestamp;
+  std::string host;
 
  private:
   void SwitchFrameInternal(bool for_top_frame);

chrome/test/chromedriver/session_commands.cc

@@ -233,6 +233,11 @@ std::unique_ptr<base::DictionaryValue> CreateCapabilities(
     caps->SetBoolean("hasTouchScreen", session->chrome->HasTouchScreen());
   }
 
+  if (session->webSocketUrl) {
+    caps->SetString("webSocketUrl",
+                    "ws://" + session->host + "/session/" + session->id);
+  }
+
   return caps;
 }
 

chrome/test/chromedriver/test/run_py_tests.py

@@ -504,6 +504,30 @@ class ChromeDriverTestWithCustomCapability(ChromeDriverBaseTestWithWebServer):
                       driver.Load, 'http://invalid/')
     self.assertEquals('http://invalid/', driver.GetCurrentUrl())
 
+class ChromeDriverWebSocketTest(ChromeDriverBaseTestWithWebServer):
+  @staticmethod
+  def composeWebSocketUrl(server_url, session_id):
+    return server_url.replace('http', 'ws') + '/session/' + session_id
+
+  def testDefaultSession(self):
+    driver = self.CreateDriver()
+    self.assertFalse(driver.capabilities.has_key('webSocketUrl'))
+
+  def testWebSocketUrlFalse(self):
+    driver = self.CreateDriver(web_socket_url=False)
+    self.assertFalse(driver.capabilities.has_key('webSocketUrl'))
+
+  def testWebSocketUrlTrue(self):
+    driver = self.CreateDriver(web_socket_url=True)
+    self.assertTrue(driver.capabilities.has_key('webSocketUrl'))
+    self.assertNotEqual(None, driver.GetSessionId())
+    self.assertEquals(driver.capabilities['webSocketUrl'],
+        self.composeWebSocketUrl(_CHROMEDRIVER_SERVER_URL,
+                                 driver.GetSessionId()))
+
+  def testWebSocketUrlInvalid(self):
+    self.assertRaises(chromedriver.InvalidArgument,
+        self.CreateDriver, web_socket_url='Invalid')
 
 class ChromeDriverTest(ChromeDriverBaseTestWithWebServer):
   """End to end tests for ChromeDriver."""