tests: Refactor to unify tests and avoid duplicated code

1 files changed, 241 insertions, 1 deletions

diff --git a/tests/default.nix b/tests/default.nix
index d9b354a..48fa9ec 100644
--- a/tests/default.nix
+++ b/tests/default.nix
@@ -3,11 +3,251 @@
   nixLib,
   pkgs,
 }: let
+  splitAttr = attr: let
+    name = builtins.elemAt (builtins.attrNames attr) 0;
+  in {
+    inherit name;
+    value = attr.${name};
+  };
+
+  vhack.runTest = {
+    name,
+    nodes ? {},
+    serverDomains ? [],
+    testScript ? {_, ...}: "",
+    services ? [],
+    ignore ? null,
+  }: let
+    # Apparently the default's (e.g. `serivces ? []`), are not applied for the `@args`
+    # syntax.
+    args = {
+      inherit name nodes serverDomains testScript services ignore;
+    };
+  in
+    specialArgs.nixos-lib.runTest {
+      hostPkgs = pkgs;
+      inherit (args) name;
+
+      meta.broken = pkgs.lib.mkIf (args.ignore != null) true;
+
+      node = {
+        specialArgs = {
+          inherit
+            (specialArgs)
+            pkgsUnstable
+            vhackPackages
+            nixpkgs-unstable
+            ;
+          inherit nixLib;
+        };
+
+        # Use the nixpkgs as constructed by the `nixpkgs.*` options
+        pkgs = null;
+      };
+
+      nodes = pkgs.lib.mkMerge [
+        {
+          acme = {
+            imports = [
+              ./common/acme/server.nix
+              ./common/dns/client.nix
+            ];
+          };
+          name_server = {nodes, ...}: {
+            imports =
+              specialArgs.extraModules
+              ++ [
+                ./common/acme/client.nix
+                ./common/dns/server.nix
+              ];
+
+            vhack.dns.zones = let
+              mkDomain = attr: let
+                split = splitAttr attr;
+              in {
+                name = "${split.value}";
+                value = {
+                  SOA = {
+                    nameServer = "ns";
+                    adminEmail = "admin@server.com";
+                    serial = 2025012301;
+                  };
+                  useOrigin = false;
+
+                  A = [
+                    nodes.${split.name}.networking.primaryIPAddress
+                  ];
+                  AAAA = [
+                    nodes.${split.name}.networking.primaryIPv6Address
+                  ];
+                };
+              };
+            in
+              builtins.listToAttrs (builtins.map mkDomain args.serverDomains);
+          };
+        }
+        args.nodes
+        (builtins.mapAttrs (_: _: {
+            imports = [
+              ./common/acme/client.nix
+              ./common/dns/client.nix
+            ];
+          })
+          args.nodes)
+      ];
+
+      testScript = {nodes, ...} @ testScriptInput: let
+        acme = import ./common/acme {inherit pkgs;};
+
+        waitFor = builtins.concatStringsSep "\n" (
+          builtins.map (attr: let
+            split = splitAttr attr;
+          in "  maybe_wait(${split.name},\"${split.value}\", ${
+            if (builtins.hasAttr "start" attr && attr.start)
+            then "True"
+            else "False"
+          })")
+          args.services
+        );
+
+        acmeOnline = builtins.concatStringsSep "\n" (builtins.map (attr: let
+          split = splitAttr attr;
+        in "    status.append(acme_online(${split.name}, \"${split.value}\", renewRan.get(\"acme-order-renew-${split.value}.service\", False)))")
+        args.serverDomains);
+
+        allRunning =
+          builtins.concatStringsSep "\n" (builtins.map (name: "  all_services_running(${name})")
+            (builtins.attrNames args.nodes));
+
+        optional = condition: value:
+          if condition
+          then value
+          else "";
+      in
+        acme.prepare (builtins.attrNames args.nodes)
+        # The acme code runs `start_all` before it hands execution back to us.
+        # So we don't need to run it.
+        (
+          optional (args.services != [] || args.serverDomains != [])
+          # Python
+          ''
+            def unit_exists(host, unitName: str) -> bool:
+              (status, _) = host.execute(f"test -f '/etc/systemd/system/{unitName}'")
+              return status == 0
+
+            wasAlive: bool
+            def run_and_wait(host, unit: str):
+              global wasAlive
+              wasAlive = False
+              def wait_state(_last_try: bool) -> bool:
+                global wasAlive
+
+                state = host.get_unit_property(unit, "SubState")
+                if state == "failed":
+                    assert False, f'Unit "{unit}" reached state "{state}"'
+
+                if state == "start" or state == "active" or state == "running":
+                    wasAlive = True
+
+                host.log(f"(waiting) {unit} -> {state}")
+                return wasAlive and state == "dead"
+
+              with host.nested(f"Waiting for unit {unit}"):
+                retry(wait_state)
+          ''
+          + optional (args.serverDomains != [])
+          # Python
+          ''
+            def acme_online(host, domain, renewRan):
+              from typing import assert_never
+
+              unitMain = f"acme-{domain}.service"
+              unitRenew = f"acme-order-renew-{domain}.service"
+
+              infoMain = host.get_unit_property(unitMain, "SubState")
+              infoRenew = host.get_unit_property(unitRenew, "SubState")
+
+              match (infoMain, infoRenew):
+                case ("exited", "dead") | ("dead", "dead") if not renewRan:
+                  # The unit might not exist, because we use the DNS values for
+                  # generating this code (there is no needed mapping between DNS -> acme
+                  # host). So we check, if it is an actual unit.
+                  if unit_exists(host, unitMain):
+                    host.start_job(f"acme-{domain}")
+                    return ("Wait", unitRenew, host)
+                  else:
+                    return ("Done", None, None)
+
+                case ("exited", "dead") | ("exited", "dead") if renewRan:
+                  # The unit is done, so it should be good?
+                  host.require_unit_state(unitMain, "active")
+                  host.succeed(f"${pkgs.lib.getExe pkgs.openssl} s_client -connect {domain}:443 -verify_return_error </dev/null")
+                  return ("Done", None, None)
+
+                case (_, "running") | (_, "start"):
+                  # It's currently running, let's wait
+                  return ("Wait", unitRenew, host)
+
+                case other:
+                  assert False, f"Wrong unit sub-states for ({unitMain},{unitRenew}): {other}"
+              assert_never() # i.e. unreachable
+
+            with subtest("Ensure, that all acme certificates have been fetched"):
+              renewRan = {}
+              while True:
+                status = []
+            ${acmeOnline}
+                done = True
+                for (state, unit, host) in status:
+                  if state == "Wait":
+                    run_and_wait(host, unit)
+                    renewRan[unit] = True
+                  if state != "Done":
+                    done = False
+                if done:
+                    break
+          ''
+          + optional (args.services != [])
+          # Python
+          ''
+            def maybe_wait(host, unit: str, should_start: bool):
+              if unit_exists(host, unit):
+                if should_start:
+                  host.start_job(unit.removesuffix(".service"))
+                  run_and_wait(host, unit)
+                else:
+                  host.wait_for_unit(unit)
+              else:
+                assert False, f"Unit '{unit}' does not exist, test misconfigured"
+
+            with subtest("Waiting for units"):
+            ${waitFor}
+          ''
+          +
+          # Python
+          ''
+            def all_services_running(host):
+              import json
+              (status, output) = host.systemctl("list-units --state=failed --plain --no-pager --output=json")
+              host_failed = json.loads(output)
+              assert len(host_failed) == 0, f"Expected zero failing services, but found: {json.dumps(host_failed, indent=4)}"
+
+            with subtest("All services running"):
+            ${allRunning}
+
+            ${(args.testScript testScriptInput)}
+
+            with subtest("All services running"):
+            ${allRunning}
+          ''
+        );
+    };
+
   tests = nixLib.mkByName {
     baseDirectory = ./by-name;
     fileName = "test.nix";
     finalizeFunction = name: value:
-      import value (nixLib.warnMerge specialArgs {inherit pkgs;} "the test args set");
+      import value (nixLib.warnMerge specialArgs {inherit pkgs vhack;} "the test args set");
   };
 in
   tests