Add a portserver_address kwarg to pick_unused_port (#12)

This lets the caller specify the address of a portserver without having
to modify the environment.

2 files changed, 32 insertions, 3 deletions

diff --git a/src/portpicker.py b/src/portpicker.py
index aee7c39..c800161 100644
--- a/src/portpicker.py
+++ b/src/portpicker.py
@@ -133,13 +133,18 @@ def is_port_free(port):
 IsPortFree = is_port_free  # legacy API. pylint: disable=invalid-name
 
 
-def pick_unused_port(pid=None):
+def pick_unused_port(pid=None, portserver_address=None):
     """A pure python implementation of PickUnusedPort.
 
     Args:
       pid: PID to tell the portserver to associate the reservation with. If
-        None,
-        the current process's PID is used.
+        None, the current process's PID is used.
+      portserver_address: The address (path) of a unix domain socket
+        with which to connect to a portserver.  A leading '@'
+        character indicates an address in the "abstract namespace." If
+        None, or no port is returned by the portserver at the provided
+        address, the environment will be checked for a PORTSERVER_ADDRESS
+        variable. If that's not set, no port server will be used.
 
     Returns:
       A port number that is unused on both TCP and UDP.
@@ -149,6 +154,10 @@ def pick_unused_port(pid=None):
         _owned_ports.add(port)
         return port
     # Provide access to the portserver on an opt-in basis.
+    if portserver_address:
+        port = get_port_from_port_server(portserver_address, pid=pid)
+        if port:
+            return port
     if 'PORTSERVER_ADDRESS' in os.environ:
         port = get_port_from_port_server(os.environ['PORTSERVER_ADDRESS'],
                                          pid=pid)
diff --git a/src/tests/portpicker_test.py b/src/tests/portpicker_test.py
index e799d01..b6ac959 100644
--- a/src/tests/portpicker_test.py
+++ b/src/tests/portpicker_test.py
@@ -70,6 +70,26 @@ class PickUnusedPortTest(unittest.TestCase):
 
     @unittest.skipIf('PORTSERVER_ADDRESS' not in os.environ,
                      'no port server to test against')
+    def testPickUnusedCanSuccessfullyUsePortServerAddressKwarg(self):
+
+        with mock.patch.object(portpicker, '_pick_unused_port_without_server'):
+            portpicker._pick_unused_port_without_server.side_effect = (
+                Exception('eek!')
+            )
+
+            # Since _PickUnusedPortWithoutServer() raises an exception, and
+            # we've temporarily removed PORTSERVER_ADDRESS from os.environ, if
+            # we can successfully obtain a port, the portserver must be working.
+            addr = os.environ.pop('PORTSERVER_ADDRESS')
+            try:
+                port = portpicker.pick_unused_port(portserver_address=addr)
+                self.assertTrue(self.IsUnusedTCPPort(port))
+                self.assertTrue(self.IsUnusedUDPPort(port))
+            finally:
+              os.environ['PORTSERVER_ADDRESS'] = addr
+
+    @unittest.skipIf('PORTSERVER_ADDRESS' not in os.environ,
+                     'no port server to test against')
     def testGetPortFromPortServer(self):
         """Exercise the get_port_from_port_server() helper function."""
         for _ in range(10):