Add CNS function to clean up port matching requester IP address.

media/tools/constrained_network_server/cns_test.py

 #!/usr/bin/env python
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Tests for Constrained Network Server."""
 import os
 import signal
 import subprocess
 import tempfile
@@ skipping 97 matching lines @@
 
     # Update fake time to see if ports expire.
     self._current_time += self._EXPIRY_TIME
 
     # Send second Get() which would normally cause ports to expire. Ensure that
     # the ports did not expire.
     self.assertEquals(self._pa.Get('test2'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)
     self.assertEquals(set(self._pa._ports.keys()), set([
         cns._DEFAULT_CNS_PORT_RANGE[0], cns._DEFAULT_CNS_PORT_RANGE[0] + 1]))
 
+  def testPortAllocatorCleanMatchingIP(self):
+    # Setup PortAllocator w/o port expiration.
+    self._pa = cns.PortAllocator(cns._DEFAULT_CNS_PORT_RANGE, 0)
+
+    # Ensure Get() succeeds and returns the correct port.
+    self.assertEquals(self._pa.Get('ip1', t=1), cns._DEFAULT_CNS_PORT_RANGE[0])
+    self.assertEquals(self._pa.Get('ip1', t=2),
+                      cns._DEFAULT_CNS_PORT_RANGE[0] + 1)
+    self.assertEquals(self._pa.Get('ip1', t=3),
+                      cns._DEFAULT_CNS_PORT_RANGE[0] + 2)
+    self.assertEquals(self._pa.Get('ip2', t=1),
+                      cns._DEFAULT_CNS_PORT_RANGE[0] + 3)
+
+    self._pa.Cleanup(all_ports=False, request_ip='ip1')
+
+    self.assertEquals(self._pa._ports.keys(),
+                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 3])
+    self.assertEquals(self._pa.Get('ip2'), cns._DEFAULT_CNS_PORT_RANGE[0])
+    self.assertEquals(self._pa.Get('ip1'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)
+
+    self._pa.Cleanup(all_ports=False, request_ip='ip2')
+    self.assertEquals(self._pa._ports.keys(),
+                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 1])
+
+    self._pa.Cleanup(all_ports=False, request_ip='abc')
+    self.assertEquals(self._pa._ports.keys(),
+                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 1])
+
+    self._pa.Cleanup(all_ports=False, request_ip='ip1')
+    self.assertEquals(self._pa._ports.keys(), [])
+
 
 class ConstrainedNetworkServerTest(unittest.TestCase):
   """End to end tests for ConstrainedNetworkServer system.
 
   These tests require root access and run the cherrypy server along with
   tc/iptables commands.
   """
 
   # Amount of time to wait for the CNS to start up.
   _SERVER_START_SLEEP_SECS = 1
@@ skipping 97 matching lines @@
 
   class DummyOptionsWithServer(object):
     www_root = 'media'
     port = 9000
     cherrypy.url = lambda: 'http://127.0.0.1:9000/ServeConstrained'
     local_server_port = 8080
 
 
 if __name__ == '__main__':
   unittest.main()