tools/telemetry/third_party/webpagereplay/trafficshaper_test.py

   1 #!/usr/bin/env python
   2 # Copyright 2011 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #      http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 """System integration test for traffic shaping.
  17
  18 Usage:
  19 $ sudo ./trafficshaper_test.py
  20 """
  21
  22 import daemonserver
  23 import logging
  24 import platformsettings
  25 import socket
  26 import SocketServer
  27 import trafficshaper
  28 import unittest
  29
  30 RESPONSE_SIZE_KEY = 'response-size:'
  31 TEST_DNS_PORT = 5555
  32 TEST_HTTP_PORT = 8888
  33 TIMER = platformsettings.timer
  34
  35
  36 def GetElapsedMs(start_time, end_time):
  37   """Return milliseconds elapsed between |start_time| and |end_time|.
  38
  39   Args:
  40     start_time: seconds as a float (or string representation of float).
  41     end_time: seconds as a float (or string representation of float).
  42   Return:
  43     milliseconds elapsed as integer.
  44   """
  45   return int((float(end_time) - float(start_time)) * 1000)
  46
  47
  48 class TrafficShaperTest(unittest.TestCase):
  49
  50   def testBadBandwidthRaises(self):
  51     self.assertRaises(trafficshaper.BandwidthValueError,
  52                       trafficshaper.TrafficShaper,
  53                       down_bandwidth='1KBit/s')
  54
  55
  56 class TimedUdpHandler(SocketServer.DatagramRequestHandler):
  57   """UDP handler that returns the time when the request was handled."""
  58
  59   def handle(self):
  60     data = self.rfile.read()
  61     read_time = self.server.timer()
  62     self.wfile.write(str(read_time))
  63
  64
  65 class TimedTcpHandler(SocketServer.StreamRequestHandler):
  66   """Tcp handler that returns the time when the request was read.
  67
  68   It can respond with the number of bytes specified in the request.
  69   The request looks like:
  70     request_data -> RESPONSE_SIZE_KEY num_reponse_bytes '\n' ANY_DATA
  71   """
  72
  73   def handle(self):
  74     data = self.rfile.read()
  75     read_time = self.server.timer()
  76     contents = str(read_time)
  77     if data.startswith(RESPONSE_SIZE_KEY):
  78       num_response_bytes = int(data[len(RESPONSE_SIZE_KEY):data.index('\n')])
  79       contents = '%s\n%s' % (contents,
  80                              '\x00' * (num_response_bytes - len(contents) - 1))
  81     self.wfile.write(contents)
  82
  83
  84 class TimedUdpServer(SocketServer.ThreadingUDPServer,
  85                      daemonserver.DaemonServer):
  86   """A simple UDP server similar to dnsproxy."""
  87
  88   # Override SocketServer.TcpServer setting to avoid intermittent errors.
  89   allow_reuse_address = True
  90
  91   def __init__(self, host, port, timer=TIMER):
  92     SocketServer.ThreadingUDPServer.__init__(
  93         self, (host, port), TimedUdpHandler)
  94     self.timer = timer
  95
  96   def cleanup(self):
  97     pass
  98
  99
 100 class TimedTcpServer(SocketServer.ThreadingTCPServer,
 101                      daemonserver.DaemonServer):
 102   """A simple TCP server similar to httpproxy."""
 103
 104   # Override SocketServer.TcpServer setting to avoid intermittent errors.
 105   allow_reuse_address = True
 106
 107   def __init__(self, host, port, timer=TIMER):
 108     SocketServer.ThreadingTCPServer.__init__(
 109         self, (host, port), TimedTcpHandler)
 110     self.timer = timer
 111
 112   def cleanup(self):
 113     try:
 114       self.shutdown()
 115     except KeyboardInterrupt, e:
 116       pass
 117
 118
 119 class TcpTestSocketCreator(object):
 120   """A TCP socket creator suitable for with-statement."""
 121
 122   def __init__(self, host, port, timeout=1.0):
 123     self.address = (host, port)
 124     self.timeout = timeout
 125
 126   def __enter__(self):
 127     self.socket = socket.create_connection(self.address, timeout=self.timeout)
 128     return self.socket
 129
 130   def __exit__(self, *args):
 131     self.socket.close()
 132
 133
 134 class TimedTestCase(unittest.TestCase):
 135   def assertValuesAlmostEqual(self, expected, actual, tolerance=0.05):
 136     """Like the following with nicer default message:
 137            assertTrue(expected <= actual + tolerance &&
 138                       expected >= actual - tolerance)
 139     """
 140     delta = tolerance * expected
 141     if actual > expected + delta or actual < expected - delta:
 142       self.fail('%s is not equal to expected %s +/- %s%%' % (
 143               actual, expected, 100 * tolerance))
 144
 145
 146 class TcpTrafficShaperTest(TimedTestCase):
 147
 148   def setUp(self):
 149     self.host = platformsettings.get_server_ip_address()
 150     self.port = TEST_HTTP_PORT
 151     self.tcp_socket_creator = TcpTestSocketCreator(self.host, self.port)
 152     self.timer = TIMER
 153
 154   def TrafficShaper(self, **kwargs):
 155     return trafficshaper.TrafficShaper(
 156         host=self.host, ports=(self.port,), **kwargs)
 157
 158   def GetTcpSendTimeMs(self, num_bytes):
 159     """Return time in milliseconds to send |num_bytes|."""
 160
 161     with self.tcp_socket_creator as s:
 162       start_time = self.timer()
 163       request_data = '\x00' * num_bytes
 164
 165       s.sendall(request_data)
 166       # TODO(slamm): Figure out why partial is shutdown needed to make it work.
 167       s.shutdown(socket.SHUT_WR)
 168       read_time = s.recv(1024)
 169     return GetElapsedMs(start_time, read_time)
 170
 171   def GetTcpReceiveTimeMs(self, num_bytes):
 172     """Return time in milliseconds to receive |num_bytes|."""
 173
 174     with self.tcp_socket_creator as s:
 175       s.sendall('%s%s\n' % (RESPONSE_SIZE_KEY, num_bytes))
 176       # TODO(slamm): Figure out why partial is shutdown needed to make it work.
 177       s.shutdown(socket.SHUT_WR)
 178       num_remaining_bytes = num_bytes
 179       read_time = None
 180       while num_remaining_bytes > 0:
 181         response_data = s.recv(4096)
 182         num_remaining_bytes -= len(response_data)
 183         if not read_time:
 184           read_time, padding = response_data.split('\n')
 185     return GetElapsedMs(read_time, self.timer())
 186
 187   def testTcpConnectToIp(self):
 188     """Verify that it takes |delay_ms| to establish a TCP connection."""
 189     if not platformsettings.has_ipfw():
 190       logging.warning('ipfw is not available in path. Skip the test')
 191       return
 192     with TimedTcpServer(self.host, self.port):
 193       for delay_ms in (100, 175):
 194         with self.TrafficShaper(delay_ms=delay_ms):
 195           start_time = self.timer()
 196           with self.tcp_socket_creator:
 197             connect_time = GetElapsedMs(start_time, self.timer())
 198         self.assertValuesAlmostEqual(delay_ms, connect_time, tolerance=0.12)
 199
 200   def testTcpUploadShaping(self):
 201     """Verify that 'up' bandwidth is shaped on TCP connections."""
 202     if not platformsettings.has_ipfw():
 203       logging.warning('ipfw is not available in path. Skip the test')
 204       return
 205     num_bytes = 1024 * 100
 206     bandwidth_kbits = 2000
 207     expected_ms = 8.0 * num_bytes / bandwidth_kbits
 208     with TimedTcpServer(self.host, self.port):
 209       with self.TrafficShaper(up_bandwidth='%sKbit/s' % bandwidth_kbits):
 210         self.assertValuesAlmostEqual(expected_ms, self.GetTcpSendTimeMs(num_bytes))
 211
 212   def testTcpDownloadShaping(self):
 213     """Verify that 'down' bandwidth is shaped on TCP connections."""
 214     if not platformsettings.has_ipfw():
 215       logging.warning('ipfw is not available in path. Skip the test')
 216       return
 217     num_bytes = 1024 * 100
 218     bandwidth_kbits = 2000
 219     expected_ms = 8.0 * num_bytes / bandwidth_kbits
 220     with TimedTcpServer(self.host, self.port):
 221       with self.TrafficShaper(down_bandwidth='%sKbit/s' % bandwidth_kbits):
 222         self.assertValuesAlmostEqual(expected_ms, self.GetTcpReceiveTimeMs(num_bytes))
 223
 224   def testTcpInterleavedDownloads(self):
 225     # TODO(slamm): write tcp interleaved downloads test
 226     pass
 227
 228
 229 class UdpTrafficShaperTest(TimedTestCase):
 230
 231   def setUp(self):
 232     self.host = platformsettings.get_server_ip_address()
 233     self.dns_port = TEST_DNS_PORT
 234     self.timer = TIMER
 235
 236   def TrafficShaper(self, **kwargs):
 237     return trafficshaper.TrafficShaper(
 238         host=self.host, ports=(self.dns_port,), **kwargs)
 239
 240   def GetUdpSendReceiveTimesMs(self):
 241     """Return time in milliseconds to send |num_bytes|."""
 242     start_time = self.timer()
 243     udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 244     udp_socket.sendto('test data\n', (self.host, self.dns_port))
 245     read_time = udp_socket.recv(1024)
 246     return (GetElapsedMs(start_time, read_time),
 247             GetElapsedMs(read_time, self.timer()))
 248
 249   def testUdpDelay(self):
 250     if not platformsettings.has_ipfw():
 251       logging.warning('ipfw is not available in path. Skip the test')
 252       return
 253     for delay_ms in (100, 170):
 254       expected_ms = delay_ms / 2
 255       with TimedUdpServer(self.host, self.dns_port):
 256         with self.TrafficShaper(delay_ms=delay_ms):
 257           send_ms, receive_ms = self.GetUdpSendReceiveTimesMs()
 258           self.assertValuesAlmostEqual(expected_ms, send_ms, tolerance=0.10)
 259           self.assertValuesAlmostEqual(expected_ms, receive_ms, tolerance=0.10)
 260
 261
 262   def testUdpInterleavedDelay(self):
 263     # TODO(slamm): write udp interleaved udp delay test
 264     pass
 265
 266
 267 class TcpAndUdpTrafficShaperTest(TimedTestCase):
 268   # TODO(slamm): Test concurrent TCP and UDP traffic
 269   pass
 270
 271
 272 # TODO(slamm): Packet loss rate (try different ports)
 273
 274
 275 if __name__ == '__main__':
 276   #logging.getLogger().setLevel(logging.DEBUG)
 277   unittest.main()