OLD | NEW |
| (Empty) |
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_response_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_byte
s)) | |
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_b
ytes)) | |
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() | |
OLD | NEW |