chrome/test/chromedriver/net/websocket_unittest.cc

   1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include <string>
   6 #include <vector>
   7
   8 #include "base/bind.h"
   9 #include "base/compiler_specific.h"
  10 #include "base/location.h"
  11 #include "base/memory/ref_counted.h"
  12 #include "base/memory/scoped_ptr.h"
  13 #include "base/message_loop/message_loop.h"
  14 #include "base/message_loop/message_loop_proxy.h"
  15 #include "base/run_loop.h"
  16 #include "base/single_thread_task_runner.h"
  17 #include "base/threading/thread.h"
  18 #include "base/time/time.h"
  19 #include "chrome/test/chromedriver/net/test_http_server.h"
  20 #include "chrome/test/chromedriver/net/websocket.h"
  21 #include "net/url_request/url_request_test_util.h"
  22 #include "testing/gtest/include/gtest/gtest.h"
  23 #include "url/gurl.h"
  24
  25 namespace {
  26
  27 void OnConnectFinished(base::RunLoop* run_loop, int* save_error, int error) {
  28   *save_error = error;
  29   run_loop->Quit();
  30 }
  31
  32 void RunPending(base::MessageLoop* loop) {
  33   base::RunLoop run_loop;
  34   loop->PostTask(FROM_HERE, run_loop.QuitClosure());
  35   run_loop.Run();
  36 }
  37
  38 class Listener : public WebSocketListener {
  39  public:
  40   explicit Listener(const std::vector<std::string>& messages)
  41       : messages_(messages) {}
  42
  43   ~Listener() override { EXPECT_TRUE(messages_.empty()); }
  44
  45   void OnMessageReceived(const std::string& message) override {
  46     ASSERT_TRUE(messages_.size());
  47     EXPECT_EQ(messages_[0], message);
  48     messages_.erase(messages_.begin());
  49     if (messages_.empty())
  50       base::MessageLoop::current()->Quit();
  51   }
  52
  53   void OnClose() override { EXPECT_TRUE(false); }
  54
  55  private:
  56   std::vector<std::string> messages_;
  57 };
  58
  59 class CloseListener : public WebSocketListener {
  60  public:
  61   explicit CloseListener(base::RunLoop* run_loop)
  62       : run_loop_(run_loop) {}
  63
  64   ~CloseListener() override { EXPECT_FALSE(run_loop_); }
  65
  66   void OnMessageReceived(const std::string& message) override {}
  67
  68   void OnClose() override {
  69     EXPECT_TRUE(run_loop_);
  70     if (run_loop_)
  71       run_loop_->Quit();
  72     run_loop_ = NULL;
  73   }
  74
  75  private:
  76   base::RunLoop* run_loop_;
  77 };
  78
  79 class WebSocketTest : public testing::Test {
  80  public:
  81   WebSocketTest() {}
  82   ~WebSocketTest() override {}
  83
  84   void SetUp() override { ASSERT_TRUE(server_.Start()); }
  85
  86   void TearDown() override { server_.Stop(); }
  87
  88  protected:
  89   scoped_ptr<WebSocket> CreateWebSocket(const GURL& url,
  90                                         WebSocketListener* listener) {
  91     int error;
  92     scoped_ptr<WebSocket> sock(new WebSocket(url, listener));
  93     base::RunLoop run_loop;
  94     sock->Connect(base::Bind(&OnConnectFinished, &run_loop, &error));
  95     loop_.PostDelayedTask(
  96         FROM_HERE, run_loop.QuitClosure(),
  97         base::TimeDelta::FromSeconds(10));
  98     run_loop.Run();
  99     if (error == net::OK)
 100       return sock.Pass();
 101     return scoped_ptr<WebSocket>();
 102   }
 103
 104   scoped_ptr<WebSocket> CreateConnectedWebSocket(WebSocketListener* listener) {
 105     return CreateWebSocket(server_.web_socket_url(), listener);
 106   }
 107
 108   void SendReceive(const std::vector<std::string>& messages) {
 109     Listener listener(messages);
 110     scoped_ptr<WebSocket> sock(CreateConnectedWebSocket(&listener));
 111     ASSERT_TRUE(sock);
 112     for (size_t i = 0; i < messages.size(); ++i) {
 113       ASSERT_TRUE(sock->Send(messages[i]));
 114     }
 115     base::RunLoop run_loop;
 116     loop_.PostDelayedTask(
 117         FROM_HERE, run_loop.QuitClosure(),
 118         base::TimeDelta::FromSeconds(10));
 119     run_loop.Run();
 120   }
 121
 122   base::MessageLoopForIO loop_;
 123   TestHttpServer server_;
 124 };
 125
 126 }  // namespace
 127
 128 TEST_F(WebSocketTest, CreateDestroy) {
 129   CloseListener listener(NULL);
 130   WebSocket sock(GURL("ws://127.0.0.1:2222"), &listener);
 131 }
 132
 133 TEST_F(WebSocketTest, Connect) {
 134   CloseListener listener(NULL);
 135   ASSERT_TRUE(CreateWebSocket(server_.web_socket_url(), &listener));
 136   RunPending(&loop_);
 137   ASSERT_TRUE(server_.WaitForConnectionsToClose());
 138 }
 139
 140 TEST_F(WebSocketTest, ConnectNoServer) {
 141   CloseListener listener(NULL);
 142   ASSERT_FALSE(CreateWebSocket(GURL("ws://127.0.0.1:33333"), NULL));
 143 }
 144
 145 TEST_F(WebSocketTest, Connect404) {
 146   server_.SetRequestAction(TestHttpServer::kNotFound);
 147   CloseListener listener(NULL);
 148   ASSERT_FALSE(CreateWebSocket(server_.web_socket_url(), NULL));
 149   RunPending(&loop_);
 150   ASSERT_TRUE(server_.WaitForConnectionsToClose());
 151 }
 152
 153 TEST_F(WebSocketTest, ConnectServerClosesConn) {
 154   server_.SetRequestAction(TestHttpServer::kClose);
 155   CloseListener listener(NULL);
 156   ASSERT_FALSE(CreateWebSocket(server_.web_socket_url(), &listener));
 157 }
 158
 159 TEST_F(WebSocketTest, CloseOnReceive) {
 160   server_.SetMessageAction(TestHttpServer::kCloseOnMessage);
 161   base::RunLoop run_loop;
 162   CloseListener listener(&run_loop);
 163   scoped_ptr<WebSocket> sock(CreateConnectedWebSocket(&listener));
 164   ASSERT_TRUE(sock);
 165   ASSERT_TRUE(sock->Send("hi"));
 166   loop_.PostDelayedTask(
 167       FROM_HERE, run_loop.QuitClosure(),
 168       base::TimeDelta::FromSeconds(10));
 169   run_loop.Run();
 170 }
 171
 172 TEST_F(WebSocketTest, CloseOnSend) {
 173   base::RunLoop run_loop;
 174   CloseListener listener(&run_loop);
 175   scoped_ptr<WebSocket> sock(CreateConnectedWebSocket(&listener));
 176   ASSERT_TRUE(sock);
 177   server_.Stop();
 178
 179   sock->Send("hi");
 180   loop_.PostDelayedTask(
 181       FROM_HERE, run_loop.QuitClosure(),
 182       base::TimeDelta::FromSeconds(10));
 183   run_loop.Run();
 184   ASSERT_FALSE(sock->Send("hi"));
 185 }
 186
 187 TEST_F(WebSocketTest, SendReceive) {
 188   std::vector<std::string> messages;
 189   messages.push_back("hello");
 190   SendReceive(messages);
 191 }
 192
 193 TEST_F(WebSocketTest, SendReceiveLarge) {
 194   std::vector<std::string> messages;
 195   messages.push_back(std::string(10 << 20, 'a'));
 196   SendReceive(messages);
 197 }
 198
 199 TEST_F(WebSocketTest, SendReceiveMultiple) {
 200   std::vector<std::string> messages;
 201   messages.push_back("1");
 202   messages.push_back("2");
 203   messages.push_back("3");
 204   SendReceive(messages);
 205 }