path: root/tests/default.nix

{
  specialArgs,
  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
            deadCounter: int
            def run_and_wait(host, unit: str):
              global wasAlive
              global deadCounter

              wasAlive = False
              deadCounter = 0

              def wait_state(_last_try: bool) -> bool:
                global wasAlive
                global deadCounter

                state = host.get_unit_property(unit, "SubState")
                host.log(f"(waiting) {unit} -> {state}")

                match state:
                  case "failed":
                    assert False, f'Unit "{unit}" reached state "{state}"'
                  case "start" | "active" | "running":
                    wasAlive = True
                  case "dead":
                    if wasAlive:
                      return True
                    else:
                      if deadCounter >= 5:
                        return True

                      deadCounter = deadCounter + 1
                      return False

                return False

              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 vhack;} "the test args set");
  };
in
  tests