tests/default: Also add an option to specify ports to wait for

1 files changed, 23 insertions, 5 deletions

diff --git a/tests/default.nix b/tests/default.nix
index b679dc1..8405e6a 100644
--- a/tests/default.nix
+++ b/tests/default.nix
@@ -42,15 +42,19 @@
     # means, that the test will not start, until `example.service` on the node with name
     # `nodeName` has started.
     services ? [],
+    # Open ports to wait for.
+    # The format follows the common one, e.g.:
+    # `{nodeName = 993;}`
+    # means, that the test will not start, until port `993` on the node with name
+    # `nodeName` is opened.
+    ports ? [],
     # If `true`, will ignore this test.
     ignore ? null,
     testScript ? {_, ...}: "",
   }: let
     # Apparently the default's (e.g. `serivces ? []`), are not applied for the `@args`
     # syntax.
-    args = {
-      inherit name dns nodes serverDomains testScript services ignore;
-    };
+    args = {inherit dns ignore name nodes ports serverDomains services testScript;};
   in
     specialArgs.nixos-lib.runTest {
       hostPkgs = pkgs;
@@ -136,7 +140,7 @@
       testScript = {nodes, ...} @ testScriptInput: let
         acme = import ./common/acme {inherit pkgs;};
 
-        waitFor = builtins.concatStringsSep "\n" (
+        waitForServices = builtins.concatStringsSep "\n" (
           builtins.map (attr: let
             split = splitAttr attr;
           in "  maybe_wait(${split.name},\"${split.value}\", ${
@@ -146,6 +150,14 @@
           })")
           args.services
         );
+        waitForPorts = builtins.concatStringsSep "\n" (
+          builtins.map (
+            attr: let
+              split = splitAttr attr;
+            in "  ${split.name}.wait_for_open_port(${builtins.toString split.value})"
+          )
+          args.ports
+        );
 
         acmeOnline = builtins.concatStringsSep "\n" (builtins.map (attr: let
           split = splitAttr attr;
@@ -274,7 +286,13 @@
                 assert False, f"Unit '{unit}' does not exist, test misconfigured"
 
             with subtest("Waiting for units"):
-            ${waitFor}
+            ${waitForServices}
+          ''
+          + optional (args.ports != [])
+          # Python
+          ''
+            with subtest("Waiting for ports"):
+            ${waitForPorts}
           ''
           +
           # Python