nixos/test-driver: add optional address arg to wait_for_{open,closed}_port

This is useful for testing servers configured to listen on specific
addresses.

3 files changed, 16 insertions, 12 deletions

diff --git a/nixos/doc/manual/development/writing-nixos-tests.section.md b/nixos/doc/manual/development/writing-nixos-tests.section.md
index 0d5bc76a2aa5..5bcdf6e58eb1 100644
--- a/nixos/doc/manual/development/writing-nixos-tests.section.md
+++ b/nixos/doc/manual/development/writing-nixos-tests.section.md
@@ -273,12 +273,13 @@ The following methods are available on machine objects:
 
 `wait_for_open_port`
 
-:   Wait until a process is listening on the given TCP port (on
-    `localhost`, at least).
+:   Wait until a process is listening on the given TCP port and IP address
+    (default `localhost`).
 
 `wait_for_closed_port`
 
-:   Wait until nobody is listening on the given TCP port.
+:   Wait until nobody is listening on the given TCP port and IP address
+    (default `localhost`).
 
 `wait_for_x`
 
diff --git a/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml b/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
index dc921dad9749..308f7c6fb0f6 100644
--- a/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
+++ b/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
@@ -483,8 +483,8 @@ start_all()
         </term>
         <listitem>
           <para>
-            Wait until a process is listening on the given TCP port (on
-            <literal>localhost</literal>, at least).
+            Wait until a process is listening on the given TCP port and
+            IP address (default <literal>localhost</literal>).
           </para>
         </listitem>
       </varlistentry>
@@ -494,7 +494,8 @@ start_all()
         </term>
         <listitem>
           <para>
-            Wait until nobody is listening on the given TCP port.
+            Wait until nobody is listening on the given TCP port and IP
+            address (default <literal>localhost</literal>).
           </para>
         </listitem>
       </varlistentry>
diff --git a/nixos/lib/test-driver/test_driver/machine.py b/nixos/lib/test-driver/test_driver/machine.py
index ffbc7c18e42b..c59ef3b17262 100644
--- a/nixos/lib/test-driver/test_driver/machine.py
+++ b/nixos/lib/test-driver/test_driver/machine.py
@@ -699,20 +699,22 @@ class Machine:
         with self.nested("waiting for file ‘{}‘".format(filename)):
             retry(check_file)
 
-    def wait_for_open_port(self, port: int) -> None:
+    def wait_for_open_port(self, port: int, addr: str = "localhost") -> None:
         def port_is_open(_: Any) -> bool:
-            status, _ = self.execute("nc -z localhost {}".format(port))
+            status, _ = self.execute("nc -z {} {}".format(addr, port))
             return status == 0
 
-        with self.nested("waiting for TCP port {}".format(port)):
+        with self.nested("waiting for TCP port {} on {}".format(port, addr)):
             retry(port_is_open)
 
-    def wait_for_closed_port(self, port: int) -> None:
+    def wait_for_closed_port(self, port: int, addr: str = "localhost") -> None:
         def port_is_closed(_: Any) -> bool:
-            status, _ = self.execute("nc -z localhost {}".format(port))
+            status, _ = self.execute("nc -z {} {}".format(addr, port))
             return status != 0
 
-        with self.nested("waiting for TCP port {} to be closed".format(port)):
+        with self.nested(
+            "waiting for TCP port {} on {} to be closed".format(port, addr)
+        ):
             retry(port_is_closed)
 
     def start_job(self, jobname: str, user: Optional[str] = None) -> Tuple[int, str]: