Merge pull request influxdata#243 from savoirfairelinux/enh_add_get_free_ports_helper_function

aviau · aviau · commit 8f128ff207eb · 2015-09-13T17:38:14.000-04:00
Enh/Fix: add a get_free_ports() "enhanced" helper function.
diff --git a/influxdb/tests/misc.py b/influxdb/tests/misc.py
@@ -1,21 +1,33 @@
-
+# -*- coding: utf-8 -*-
 
 import socket
-import time
 
 
-def get_free_port(ip='127.0.0.1'):
-    sock = socket.socket()
+def get_free_ports(num_ports, ip='127.0.0.1'):
+    """Get `num_ports` free/available ports on the interface linked to the `ip´
+    :param int num_ports: The number of free ports to get
+    :param str ip: The ip on which the ports have to be taken
+    :return: a set of ports number
+    """
+    sock_ports = []
+    ports = set()
     try:
-        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        sock.bind((ip, 0))
-        return sock.getsockname()[1]
+        for _ in range(num_ports):
+            sock = socket.socket()
+            cur = [sock, -1]
+            # append the socket directly,
+            # so that it'll be also closed (no leaked resource)
+            # in the finally here after.
+            sock_ports.append(cur)
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind((ip, 0))
+            cur[1] = sock.getsockname()[1]
     finally:
-        sock.close()
-
-        # Is there a better way than a sleep?
-        # There were issues on Travis where the port was not yet free.
-        time.sleep(0.1)
+        for sock, port in sock_ports:
+            sock.close()
+            ports.add(port)
+    assert num_ports == len(ports)
+    return ports
 
 
 def is_port_open(port, ip='127.0.0.1'):
diff --git a/influxdb/tests/server_tests/influxdb_instance.py b/influxdb/tests/server_tests/influxdb_instance.py
@@ -12,7 +12,7 @@
 import unittest
 import sys
 
-from influxdb.tests.misc import get_free_port, is_port_open
+from influxdb.tests.misc import is_port_open, get_free_ports
 
 
 class InfluxDbInstance(object):
@@ -54,12 +54,12 @@ def _start_server(self, conf_template, udp_enabled):
             dir=self.temp_dir_base)
 
         # find a couple free ports :
-        ports = dict(
-            http_port=get_free_port(),
-            admin_port=get_free_port(),
-            meta_port=get_free_port(),
-            udp_port=get_free_port() if udp_enabled else -1,
-        )
+        free_ports = get_free_ports(4)
+        ports = {}
+        for service in 'http', 'admin', 'meta', 'udp':
+            ports[service + '_port'] = free_ports.pop()
+        if not udp_enabled:
+            ports['udp_port'] = -1
 
         conf_data = dict(
             meta_dir=os.path.join(tempdir, 'meta'),
