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