src/tools/telemetry/telemetry/core/backends/chrome/websocket_unittest.py

   1 # Copyright 2013 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 import base64
   6 import BaseHTTPServer
   7 import hashlib
   8 import socket
   9 import threading
  10 import unittest
  11
  12
  13 from telemetry.core.backends.chrome import websocket
  14
  15
  16 # Minimal handler for a local websocket server.
  17 class _FakeWebSocketHandler(BaseHTTPServer.BaseHTTPRequestHandler):
  18   def do_GET(self):
  19     key = self.headers.getheader('Sec-WebSocket-Key')
  20
  21     value = key + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
  22     hashed = base64.encodestring(hashlib.sha1(value).digest()).strip().lower()
  23
  24     self.send_response(101)
  25
  26     self.send_header('Sec-Websocket-Accept', hashed)
  27     self.send_header('upgrade', 'websocket')
  28     self.send_header('connection', 'upgrade')
  29     self.end_headers()
  30
  31     self.wfile.flush()
  32
  33
  34 class TestWebSocket(unittest.TestCase):
  35   def testExports(self):
  36     self.assertNotEqual(websocket.create_connection, None)
  37     self.assertNotEqual(websocket.WebSocketException, None)
  38     self.assertNotEqual(websocket.WebSocketTimeoutException, None)
  39
  40   def testSockOpts(self):
  41     httpd = BaseHTTPServer.HTTPServer(('', 0), _FakeWebSocketHandler)
  42     threading.Thread(target=httpd.handle_request).start()
  43     ws_url = 'ws://127.0.0.1:%d' % httpd.server_port
  44     ws = websocket.create_connection(ws_url)
  45     self.assertNotEquals(
  46         ws.sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR), 0)
  47
  48     threading.Thread(target=httpd.handle_request).start()
  49     ws = websocket.create_connection(
  50         ws_url,
  51         sockopt=[(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)])
  52     self.assertNotEquals(
  53         ws.sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR), 0)
  54     self.assertNotEquals(
  55         ws.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0)