use 'port-for' for port generation

Author: funbringer

Dockerfile.tmpl

@@ -3,11 +3,11 @@ FROM postgres:${PG_VERSION}-alpine
 ENV PYTHON=python${PYTHON_VERSION}
 RUN if [ "${PYTHON_VERSION}" = "2" ] ; then \
 	apk add --no-cache python py-pip; \
-	pip install six pg8000 flake8; \
+	pip install six pg8000 flake8 port-for; \
 	fi
 RUN if [ "${PYTHON_VERSION}" = "3" ] ; then \
 	apk add --no-cache python3; \
-	pip3 install six pg8000 flake8; \
+	pip3 install six pg8000 flake8 port-for; \
 	fi
 ENV LANG=C.UTF-8
 

setup.py

@@ -2,13 +2,13 @@
 setup(
     name='testgres',
     packages=['testgres'],
-    version='0.3.1',
+    version='0.4.0',
     description='Testing utility for postgresql and its extensions',
     license='PostgreSQL',
     author='Ildar Musin',
     author_email='zildermann@gmail.com',
     url='https://github.com/postgrespro/testgres',
     keywords=['testing', 'postgresql'],
     classifiers=[],
-    install_requires=["pg8000", "six"]
+    install_requires=["pg8000", "six", "port-for"]
 )

testgres/testgres.py

@@ -22,14 +22,13 @@
 """
 
 import os
-import random
-# import socket
 import subprocess
 import pwd
 import tempfile
 import shutil
 import time
 import six
+import port_for
 
 import threading
 import logging
@@ -48,10 +47,10 @@
         raise ImportError("You must have psycopg2 or pg8000 modules installed")
 
 
+bound_ports = set()
 registered_nodes = []
 util_threads = []
 tmp_dirs = []
-last_assigned_port = int(random.random() * 16384) + 49152
 pg_config_data = {}
 base_data_dir = None
 
@@ -72,6 +71,14 @@ class QueryException(Exception):
     pass
 
 
+class InitPostgresNodeException(Exception):
+
+    """
+    Predefined exceptions
+    """
+    pass
+
+
 class PgcontroldataException(Exception):
     def __init__(self, error_text, cmd):
         self.error_text = error_text
@@ -205,6 +212,16 @@ def close(self):
 class PostgresNode(object):
 
     def __init__(self, name, port, base_dir=None, use_logging=False):
+        global bound_ports
+
+        # check that port is not used
+        if port in bound_ports:
+            raise InitPostgresNodeException(
+                    'port {} is already in use'.format(port))
+
+        # mark port as used
+        bound_ports.add(port)
+
         self.name = name
         self.host = '127.0.0.1'
         self.port = port
@@ -220,6 +237,18 @@ def __init__(self, name, port, base_dir=None, use_logging=False):
         self.use_logging = use_logging
         self.logger = None
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        global bound_ports
+
+        # stop node if necessary
+        self.cleanup()
+
+        # mark port as free
+        bound_ports.remove(self.port)
+
     @property
     def data_dir(self):
         return os.path.join(self.base_dir, "data")
@@ -718,31 +747,12 @@ def version_to_num(version):
 
 def get_new_node(name, base_dir=None, use_logging=False):
     global registered_nodes
-    global last_assigned_port
-
-    port = last_assigned_port + 1
-
-    # TODO: check if port is already used
-
-    # found = False
-    # while found:
-    #   # Check first that candidate port number is not included in
-    #   # the list of already-registered nodes.
-    #   found = True
-    #   for node in registered_nodes:
-    #       if node.port == port:
-    #           found = False
-    #           break
 
-    #   if found:
-    #       socket(socket.SOCK,
-    #              socket.PF_INET,
-    #              socket.SOCK_STREAM,
-    #              socket.getprotobyname("tcp"))
+    # generate a new unique port
+    port = port_for.select_random(exclude_ports=bound_ports)
 
     node = PostgresNode(name, port, base_dir, use_logging=use_logging)
     registered_nodes.append(node)
-    last_assigned_port = port
 
     return node
 

testgres/tests/test_simple.py

@@ -6,7 +6,7 @@
 import tempfile
 import logging.config
 
-from testgres import get_new_node, stop_all, get_config, clean_all
+from testgres import get_new_node, stop_all, get_config, clean_all, bound_ports
 
 
 class SimpleTest(unittest.TestCase):
@@ -17,79 +17,70 @@ def teardown(self):
 
     @unittest.skip("demo")
     def test_start_stop(self):
-        node = get_new_node('test')
-        node.init()
-        node.start()
-        res = node.execute('postgres', 'select 1')
-        self.assertEqual(len(res), 1)
-        self.assertEqual(res[0][0], 1)
-        node.stop()
+        with get_new_node('test') as node:
+            node.init().start()
+            res = node.execute('postgres', 'select 1')
+            self.assertEqual(len(res), 1)
+            self.assertEqual(res[0][0], 1)
 
     def test_backup_and_replication(self):
-        node = get_new_node('test')
-        replica = get_new_node('repl')
-
-        node.init(allows_streaming=True)
-        node.start()
-        node.psql('postgres', 'create table abc(a int, b int)')
-        node.psql('postgres', 'insert into abc values (1, 2)')
-        node.backup('my_backup')
-
-        replica.init_from_backup(node, 'my_backup', has_streaming=True)
-        replica.start()
-        res = replica.execute('postgres', 'select * from abc')
-        self.assertEqual(len(res), 1)
-        self.assertEqual(res[0][0], 1)
-        self.assertEqual(res[0][1], 2)
-
-        # Prepare the query which would check whether record reached replica
-        # (It is slightly different for Postgres 9.6 and Postgres 10+)
-        if get_config()['VERSION_NUM'] >= 1000000:
-            wait_lsn = 'SELECT pg_current_wal_lsn() <= replay_lsn '           \
-                'FROM pg_stat_replication WHERE application_name = \'%s\''    \
-                % replica.name
-        else:
-            wait_lsn = 'SELECT pg_current_xlog_location() <= replay_location '\
-                'FROM pg_stat_replication WHERE application_name = \'%s\''    \
-                % replica.name
-
-        # Insert into master node
-        node.psql('postgres', 'insert into abc values (3, 4)')
-        # Wait until data syncronizes
-        node.poll_query_until('postgres', wait_lsn)
-        # Check that this record was exported to replica
-        res = replica.execute('postgres', 'select * from abc')
-        self.assertEqual(len(res), 2)
-        self.assertEqual(res[1][0], 3)
-        self.assertEqual(res[1][1], 4)
-
-        node.stop()
-        replica.stop()
+        with get_new_node('test') as node, get_new_node('repl') as replica:
+            node.init(allows_streaming=True)
+            node.start()
+            node.psql('postgres', 'create table abc(a int, b int)')
+            node.psql('postgres', 'insert into abc values (1, 2)')
+            node.backup('my_backup')
+
+            replica.init_from_backup(node, 'my_backup', has_streaming=True)
+            replica.start()
+            res = replica.execute('postgres', 'select * from abc')
+            self.assertEqual(len(res), 1)
+            self.assertEqual(res[0][0], 1)
+            self.assertEqual(res[0][1], 2)
+
+            # Prepare the query which would check whether record reached replica
+            # (It is slightly different for Postgres 9.6 and Postgres 10+)
+            if get_config()['VERSION_NUM'] >= 1000000:
+                wait_lsn = 'SELECT pg_current_wal_lsn() <= replay_lsn '           \
+                    'FROM pg_stat_replication WHERE application_name = \'%s\''    \
+                    % replica.name
+            else:
+                wait_lsn = 'SELECT pg_current_xlog_location() <= replay_location '\
+                    'FROM pg_stat_replication WHERE application_name = \'%s\''    \
+                    % replica.name
+
+            # Insert into master node
+            node.psql('postgres', 'insert into abc values (3, 4)')
+            # Wait until data syncronizes
+            node.poll_query_until('postgres', wait_lsn)
+            # Check that this record was exported to replica
+            res = replica.execute('postgres', 'select * from abc')
+            self.assertEqual(len(res), 2)
+            self.assertEqual(res[1][0], 3)
+            self.assertEqual(res[1][1], 4)
 
     def test_dump(self):
-        node = get_new_node('test')
-        node.init().start()
-        node.safe_psql(
-            'postgres',
-            'create table abc as '
-            'select g as a, g as b from generate_series(1, 10) as g'
-        )
-        node.psql('postgres', 'create database test')
-        node.dump('postgres', 'test.sql')
-        node.restore('test', 'test.sql')
-        self.assertEqual(
-            node.psql('postgres', 'select * from abc'),
-            node.psql('test', 'select * from abc'),
-        )
-        node.stop()
+        with get_new_node('test') as node:
+            node.init().start()
+            node.safe_psql(
+                'postgres',
+                'create table abc as '
+                'select g as a, g as b from generate_series(1, 10) as g'
+            )
+            node.psql('postgres', 'create database test')
+            node.dump('postgres', 'test.sql')
+            node.restore('test', 'test.sql')
+            self.assertEqual(
+                node.psql('postgres', 'select * from abc'),
+                node.psql('test', 'select * from abc'),
+            )
 
     def test_users(self):
-        node = get_new_node('master')
-        node.init().start()
-        node.psql('postgres', 'create role test_user login')
-        value = node.safe_psql('postgres', 'select 1', username='test_user')
-        self.assertEqual(value, six.b('1\n'))
-        node.stop()
+        with get_new_node('master') as node:
+            node.init().start()
+            node.psql('postgres', 'create role test_user login')
+            value = node.safe_psql('postgres', 'select 1', username='test_user')
+            self.assertEqual(value, six.b('1\n'))
 
     def test_logging(self):
         regex = re.compile('.+?LOG:.*')
@@ -118,21 +109,33 @@ def test_logging(self):
 
         logging.config.dictConfig(log_conf)
 
-        node = get_new_node('master', use_logging=True)
-        node1 = get_new_node('slave1', use_logging=True)
-        node2 = get_new_node('slave2', use_logging=True)
+        with get_new_node('master', use_logging=True) as node, \
+            get_new_node('slave1', use_logging=True) as node1, \
+            get_new_node('slave2', use_logging=True) as node2:
 
-        node.init().start()
-        node1.init().start()
-        node2.init().start()
+            node.init().start()
+            node1.init().start()
+            node2.init().start()
 
-        with open(logfile.name, 'r') as log:
-            for line in log:
-                self.assertTrue(regex.match(line))
+            with open(logfile.name, 'r') as log:
+                for line in log:
+                    self.assertTrue(regex.match(line))
 
-        node.stop()
-        node1.stop()
-        node2.stop()
+    def test_ports_management(self):
+        # check that no ports have been bound yet
+        self.assertEqual(len(bound_ports), 0)
+
+        with get_new_node('node') as node:
+            # check that we've just bound a port
+            self.assertEqual(len(bound_ports), 1)
+
+            # check that bound_ports contains our port
+            port_1 = list(bound_ports)[0]
+            port_2 = node.port
+            self.assertEqual(port_1, port_2)
+
+        # check that port has been freed successfuly
+        self.assertEqual(len(bound_ports), 0)
 
 
 if __name__ == '__main__':