1#!/usr/bin/env python 2 3# Copyright (c) 2012 The Chromium Authors. All rights reserved. 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7"""Tests for traffic control library.""" 8import unittest 9 10import traffic_control 11 12 13class TrafficControlUnitTests(unittest.TestCase): 14 """Unit tests for traffic control.""" 15 16 # Stores commands called by the traffic control _Exec function. 17 commands = [] 18 19 def _ExecMock(self, command, **kwargs): 20 """Mocks traffic_control._Exec and adds the command to commands list.""" 21 cmd_list = [str(x) for x in command] 22 self.commands.append(' '.join(cmd_list)) 23 return '' 24 25 def setUp(self): 26 """Resets the commands list and set the _Exec mock function.""" 27 self.commands = [] 28 self._old_Exec = traffic_control._Exec 29 traffic_control._Exec = self._ExecMock 30 31 def tearDown(self): 32 """Resets the _Exec mock function to the original.""" 33 traffic_control._Exec = self._old_Exec 34 35 def testCreateConstrainedPort(self): 36 config = { 37 'interface': 'fakeeth', 38 'port': 12345, 39 'server_port': 8888, 40 'bandwidth': 256, 41 'latency': 100, 42 'loss': 2 43 } 44 traffic_control.CreateConstrainedPort(config) 45 expected = [ 46 'sudo tc qdisc add dev fakeeth root handle 1: htb', 47 'sudo tc class add dev fakeeth parent 1: classid 1:3039 htb rate ' 48 '256kbit ceil 256kbit', 49 'sudo tc qdisc add dev fakeeth parent 1:3039 handle 3039:0 netem loss ' 50 '2% delay 100ms', 51 'sudo tc filter add dev fakeeth protocol ip parent 1: prio 1 u32 match ' 52 'ip sport 12345 0xffff flowid 1:3039', 53 'sudo iptables -t nat -A PREROUTING -i fakeeth -p tcp --dport 12345 -j ' 54 'REDIRECT --to-port 8888', 55 'sudo iptables -t nat -A OUTPUT -p tcp --dport 12345 -j REDIRECT ' 56 '--to-port 8888' 57 ] 58 self.assertEqual(expected, self.commands) 59 60 def testCreateConstrainedPortDefaults(self): 61 config = { 62 'interface': 'fakeeth', 63 'port': 12345, 64 'server_port': 8888, 65 'latency': None 66 } 67 traffic_control.CreateConstrainedPort(config) 68 expected = [ 69 'sudo tc qdisc add dev fakeeth root handle 1: htb', 70 'sudo tc class add dev fakeeth parent 1: classid 1:3039 htb rate ' 71 '%dkbit ceil %dkbit' % (traffic_control._DEFAULT_MAX_BANDWIDTH_KBIT, 72 traffic_control._DEFAULT_MAX_BANDWIDTH_KBIT), 73 'sudo tc qdisc add dev fakeeth parent 1:3039 handle 3039:0 netem', 74 'sudo tc filter add dev fakeeth protocol ip parent 1: prio 1 u32 ' 75 'match ip sport 12345 0xffff flowid 1:3039', 76 'sudo iptables -t nat -A PREROUTING -i fakeeth -p tcp --dport 12345 -j ' 77 'REDIRECT --to-port 8888', 78 'sudo iptables -t nat -A OUTPUT -p tcp --dport 12345 -j REDIRECT ' 79 '--to-port 8888' 80 ] 81 self.assertEqual(expected, self.commands) 82 83 def testDeleteConstrainedPort(self): 84 config = { 85 'interface': 'fakeeth', 86 'port': 12345, 87 'server_port': 8888, 88 'bandwidth': 256, 89 } 90 _old_GetFilterHandleId = traffic_control._GetFilterHandleId 91 traffic_control._GetFilterHandleId = lambda interface, port: '800::800' 92 93 try: 94 traffic_control.DeleteConstrainedPort(config) 95 expected = [ 96 'sudo tc filter del dev fakeeth protocol ip parent 1:0 handle ' 97 '800::800 prio 1 u32', 98 'sudo tc class del dev fakeeth parent 1: classid 1:3039 htb rate ' 99 '256kbit ceil 256kbit', 100 'sudo iptables -t nat -D PREROUTING -i fakeeth -p tcp --dport 12345 ' 101 '-j REDIRECT --to-port 8888', 102 'sudo iptables -t nat -D OUTPUT -p tcp --dport 12345 -j REDIRECT ' 103 '--to-port 8888'] 104 self.assertEqual(expected, self.commands) 105 finally: 106 traffic_control._GetFilterHandleId = _old_GetFilterHandleId 107 108 def testTearDown(self): 109 config = {'interface': 'fakeeth'} 110 111 traffic_control.TearDown(config) 112 expected = [ 113 'sudo tc qdisc del dev fakeeth root', 114 'sudo iptables -t nat -F' 115 ] 116 self.assertEqual(expected, self.commands) 117 118 def testGetFilterHandleID(self): 119 # Check seach for handle ID command. 120 self.assertRaises(traffic_control.TrafficControlError, 121 traffic_control._GetFilterHandleId, 'fakeeth', 1) 122 self.assertEquals(self.commands, ['sudo tc filter list dev fakeeth parent ' 123 '1:']) 124 125 # Check with handle ID available. 126 traffic_control._Exec = (lambda command, msg: 127 'filter parent 1: protocol ip pref 1 u32 fh 800::800 order 2048 key ht ' 128 '800 bkt 0 flowid 1:1\nmatch 08ae0000/ffff0000 at 20') 129 output = traffic_control._GetFilterHandleId('fakeeth', 1) 130 self.assertEqual(output, '800::800') 131 132 # Check with handle ID not available. 133 traffic_control._Exec = (lambda command, msg: 134 'filter parent 1: protocol ip pref 1 u32 fh 800::800 order 2048 key ht ' 135 '800 bkt 0 flowid 1:11\nmatch 08ae0000/ffff0000 at 20') 136 self.assertRaises(traffic_control.TrafficControlError, 137 traffic_control._GetFilterHandleId, 'fakeeth', 1) 138 139 traffic_control._Exec = lambda command, msg: 'NO ID IN HERE' 140 self.assertRaises(traffic_control.TrafficControlError, 141 traffic_control._GetFilterHandleId, 'fakeeth', 1) 142 143 144if __name__ == '__main__': 145 unittest.main() 146