improved socket related tests

test/wait-for-it.py

@@ -23,6 +23,8 @@ except:
 MISSING_ARGS_TEXT = "Error: you need to provide a host and port to test."
 HELP_TEXT = "Usage:"  # Start of help text
 DIVIDE_LINE = '-' * 71  # Output line of dashes
+_SOCKET_TIMEOUT = 5.0
+socket.setdefaulttimeout(_SOCKET_TIMEOUT)
 
 
 class TestWaitForIt(unittest.TestCase):
@@ -51,10 +53,13 @@ class TestWaitForIt(unittest.TestCase):
             returncode = proc.poll()
         return (returncode, _stdout, _stderr)
 
-    def open_local_port(self, host="localhost", port=8929, timeout=5):
-        s = socket.socket()
+    def open_local_port(self, host="localhost", port=8929, timeout=None):
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        if timeout:
+            s.settimeout(timeout)
         s.bind((host, port))
-        s.listen(timeout)
+        s.listen(1)
         return s
 
     def check_args(self, args, stdout_regex, stderr_regex, exitcode):
@@ -103,23 +108,27 @@ class TestWaitForIt(unittest.TestCase):
 
     def test_host_port(self):
         """Check that --host and --port args work correctly."""
-        soc = self.open_local_port(port=8929)
-        self.check_args(
-            "--host=localhost --port=8929 --timeout=1", "",
-            "wait-for-it.sh: waiting 1 seconds for localhost:8929", 0)
-        soc.close()
+        try:
+            soc = self.open_local_port(port=8929)
+            self.check_args(
+                "--host=localhost --port=8929 --timeout=1", "",
+                "wait-for-it.sh: waiting 1 seconds for localhost:8929", 0)
+        finally:
+            soc.close()
 
     def test_combined_host_port(self):
         """Tests that wait-for-it.sh returns correctly after establishing a
         connectionm using combined host and ports."""
-        soc = self.open_local_port(port=8929)
-        self.check_args(
-            "localhost:8929 --timeout=1", "",
-            "wait-for-it.sh: waiting 1 seconds for localhost:8929", 0)
-        soc.close()
+        try:
+            soc = self.open_local_port(port=8929)
+            self.check_args(
+                "localhost:8929 --timeout=1", "",
+                "wait-for-it.sh: waiting 1 seconds for localhost:8929", 0)
+        finally:
+            soc.close()
 
     def test_port_failure_with_timeout(self):
-        """Note exit status of 124 is exected, passed from the timeout
+        """Note exit status of 124 is expected, passed from the timeout
         command."""
         self.check_args(
             "localhost:8929 --timeout=1", "",
@@ -129,21 +138,27 @@ class TestWaitForIt(unittest.TestCase):
     def test_command_execution(self):
         """Checks that a command executes correctly after a port test
         passes."""
-        soc = self.open_local_port(port=8929)
-        self.check_args(
-            "localhost:8929 -- echo \"CMD OUTPUT\"", "CMD OUTPUT",
-            ".*wait-for-it.sh: localhost:8929 is available after 0 seconds", 0)
-        soc.close()
+        try:
+            soc = self.open_local_port(port=8929)
+            self.check_args(
+                "localhost:8929 -- echo \"CMD OUTPUT\"", "CMD OUTPUT",
+                ".*wait-for-it.sh: localhost:8929 is available after 0 seconds",
+                0)
+        finally:
+            soc.close()
 
     def test_failed_command_execution(self):
         """Check command failure.
 
         The command in question outputs STDERR and an exit code of 2
         """
-        soc = self.open_local_port(port=8929)
-        self.check_args("localhost:8929 -- ls not_real_file", "",
-                        ".*No such file or directory\n", 2)
-        soc.close()
+        expected_exitcode, _, _ = self.execute('ls not_real_file')
+        try:
+            soc = self.open_local_port(port=8929)
+            self.check_args("localhost:8929 -- ls not_real_file", "",
+                            ".*No such file or directory\n", expected_exitcode)
+        finally:
+            soc.close()
 
     def test_command_after_connection_failure(self):
         """Test that a command still runs even if a connection times out and