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