• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2# Copyright 2012 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"""Unit tests for proxyshaper.
17
18Usage:
19$ ./proxyshaper_test.py
20"""
21
22import proxyshaper
23import StringIO
24import unittest
25
26
27# pylint: disable=bad-whitespace
28VALID_RATES = (
29    # input,       expected_bps
30    ( '384Kbit/s',   384000),
31    ('1536Kbit/s',  1536000),
32    (   '1Mbit/s',  1000000),
33    (   '5Mbit/s',  5000000),
34    (  '2MByte/s', 16000000),
35    (         '0',        0),
36    (         '5',        5),
37    (      384000,   384000),
38    )
39
40ERROR_RATES = (
41    '1536KBit/s',  # Older versions of dummynet used capital 'B' for bytes.
42    '1Mbyte/s',    # Require capital 'B' for bytes.
43    '5bps',
44    )
45
46
47class TimedTestCase(unittest.TestCase):
48  def assertValuesAlmostEqual(self, expected, actual, tolerance=0.05):
49    """Like the following with nicer default message:
50           assertTrue(expected <= actual + tolerance &&
51                      expected >= actual - tolerance)
52    """
53    delta = tolerance * expected
54    if actual > expected + delta or actual < expected - delta:
55      self.fail('%s is not equal to expected %s +/- %s%%' % (
56              actual, expected, 100 * tolerance))
57
58
59class RateLimitedFileTest(TimedTestCase):
60  def testReadLimitedBasic(self):
61    num_bytes = 1024
62    bps = 384000
63    request_counter = lambda: 1
64    f = StringIO.StringIO(' ' * num_bytes)
65    limited_f = proxyshaper.RateLimitedFile(request_counter, f, bps)
66    start = proxyshaper.TIMER()
67    self.assertEqual(num_bytes, len(limited_f.read()))
68    expected_ms = 8.0 * num_bytes / bps * 1000.0
69    actual_ms = (proxyshaper.TIMER() - start) * 1000.0
70    self.assertValuesAlmostEqual(expected_ms, actual_ms)
71
72  def testReadlineLimitedBasic(self):
73    num_bytes = 1024 * 8 + 512
74    bps = 384000
75    request_counter = lambda: 1
76    f = StringIO.StringIO(' ' * num_bytes)
77    limited_f = proxyshaper.RateLimitedFile(request_counter, f, bps)
78    start = proxyshaper.TIMER()
79    self.assertEqual(num_bytes, len(limited_f.readline()))
80    expected_ms = 8.0 * num_bytes / bps * 1000.0
81    actual_ms = (proxyshaper.TIMER() - start) * 1000.0
82    self.assertValuesAlmostEqual(expected_ms, actual_ms)
83
84  def testReadLimitedSlowedByMultipleRequests(self):
85    num_bytes = 1024
86    bps = 384000
87    request_count = 2
88    request_counter = lambda: request_count
89    f = StringIO.StringIO(' ' * num_bytes)
90    limited_f = proxyshaper.RateLimitedFile(request_counter, f, bps)
91    start = proxyshaper.TIMER()
92    num_read_bytes = limited_f.read()
93    self.assertEqual(num_bytes, len(num_read_bytes))
94    expected_ms = 8.0 * num_bytes / (bps / float(request_count)) * 1000.0
95    actual_ms = (proxyshaper.TIMER() - start) * 1000.0
96    self.assertValuesAlmostEqual(expected_ms, actual_ms)
97
98  def testWriteLimitedBasic(self):
99    num_bytes = 1024 * 10 + 350
100    bps = 384000
101    request_counter = lambda: 1
102    f = StringIO.StringIO()
103    limited_f = proxyshaper.RateLimitedFile(request_counter, f, bps)
104    start = proxyshaper.TIMER()
105    limited_f.write(' ' * num_bytes)
106    self.assertEqual(num_bytes, len(limited_f.getvalue()))
107    expected_ms = 8.0 * num_bytes / bps * 1000.0
108    actual_ms = (proxyshaper.TIMER() - start) * 1000.0
109    self.assertValuesAlmostEqual(expected_ms, actual_ms)
110
111  def testWriteLimitedSlowedByMultipleRequests(self):
112    num_bytes = 1024 * 10
113    bps = 384000
114    request_count = 2
115    request_counter = lambda: request_count
116    f = StringIO.StringIO(' ' * num_bytes)
117    limited_f = proxyshaper.RateLimitedFile(request_counter, f, bps)
118    start = proxyshaper.TIMER()
119    limited_f.write(' ' * num_bytes)
120    self.assertEqual(num_bytes, len(limited_f.getvalue()))
121    expected_ms = 8.0 * num_bytes / (bps / float(request_count)) * 1000.0
122    actual_ms = (proxyshaper.TIMER() - start) * 1000.0
123    self.assertValuesAlmostEqual(expected_ms, actual_ms)
124
125
126class GetBitsPerSecondTest(unittest.TestCase):
127  def testConvertsValidValues(self):
128    for dummynet_option, expected_bps in VALID_RATES:
129      bps = proxyshaper.GetBitsPerSecond(dummynet_option)
130      self.assertEqual(
131          expected_bps, bps, 'Unexpected result for %s: %s != %s' % (
132              dummynet_option, expected_bps, bps))
133
134  def testRaisesOnUnexpectedValues(self):
135    for dummynet_option in ERROR_RATES:
136      self.assertRaises(proxyshaper.BandwidthValueError,
137                        proxyshaper.GetBitsPerSecond, dummynet_option)
138
139
140if __name__ == '__main__':
141  unittest.main()
142