integration-tests: Start migration to modular test framework

Still need to migrate most logic in tools.nix to modules.

integration-tests/allow-apply-all/default.nix

@@ -4,9 +4,10 @@ let
   tools = pkgs.callPackage ../tools.nix {
     targets = [ "alpha" ];
   };
-in tools.makeTest {
+in tools.runTest {
   name = "colmena-allow-apply-all";
 
+  colmena.test = {
     bundle = ./.;
 
     testScript = ''
@@ -16,4 +17,5 @@ in tools.makeTest {
 
       deployer.succeed("cd /tmp/bundle && run-copy-stderr ${tools.colmenaExec} apply --on @target")
     '';
+  };
 }

integration-tests/apply-local/default.nix

@@ -12,9 +12,10 @@ let
       security.sudo.wheelNeedsPassword = false;
     };
   };
-in tools.makeTest {
+in tools.runTest {
   name = "colmena-apply-local";
 
+  colmena.test = {
     bundle = ./.;
 
     testScript = ''
@@ -22,4 +23,5 @@ in tools.makeTest {
       deployer.succeed("grep SUCCESS /etc/deployment")
       deployer.succeed("grep SECRET /run/keys/key-text")
     '';
+  };
 }

integration-tests/apply/default.nix

@@ -4,13 +4,14 @@
 
 let
   tools = pkgs.callPackage ../tools.nix {};
-in tools.makeTest {
+in tools.runTest {
   name = "colmena-apply-${evaluator}";
 
+  colmena.test = {
     bundle = ./.;
-
     testScript = ''
       colmena = "${tools.colmenaExec}"
       evaluator = "${evaluator}"
     '' + builtins.readFile ./test-script.py;
+  };
 }

integration-tests/build-on-target/default.nix

@@ -5,9 +5,10 @@ let
     deployers = [ "deployer" "alpha" "beta" ];
     targets = [];
   };
-in tools.makeTest {
+in tools.runTest {
   name = "colmena-build-on-target";
 
+  colmena.test = {
     bundle = ./.;
 
     testScript = ''
@@ -32,4 +33,5 @@ in tools.makeTest {
           profile = beta.succeed("readlink /run/current-system")
           deployer.fail(f"nix-store -qR {profile}")
     '';
+  };
 }

integration-tests/exec/default.nix

@@ -2,9 +2,10 @@
 
 let
   tools = pkgs.callPackage ../tools.nix {};
-in tools.makeTest {
+in tools.runTest {
   name = "colmena-exec";
 
+  colmena.test = {
     bundle = ./.;
 
     testScript = ''
@@ -14,4 +15,5 @@ in tools.makeTest {
       assert "output from beta" in logs
       assert "output from gamma" in logs
     '';
+  };
 }

integration-tests/flakes/default.nix

@@ -6,9 +6,10 @@ let
   tools = pkgs.callPackage ../tools.nix {
     targets = [ "alpha" ];
   };
-in tools.makeTest {
+in tools.runTest {
   name = "colmena-flakes-${evaluator}";
 
+  colmena.test = {
     bundle = ./.;
 
     testScript = ''
@@ -45,4 +46,5 @@ in tools.makeTest {
           deployer.succeed("cd /tmp/bundle && ${tools.colmenaExec} apply --on @target --evaluator ${evaluator} --impure")
           alpha.succeed("grep deployer /etc/deployment")
     '';
+  };
 }

integration-tests/parallel/default.nix

@@ -5,6 +5,7 @@ let
 in tools.makeTest {
   name = "colmena-parallel";
 
+  colmena.test = {
     bundle = ./.;
 
     testScript = ''
@@ -26,4 +27,5 @@ in tools.makeTest {
 
           assert delay < 2000
     '';
+  };
 }

integration-tests/tools.nix

@@ -3,9 +3,7 @@
 # By default, we have four nodes: deployer, alpha, beta, gamma.
 # deployer is where colmena will run.
 #
-# `nixos/lib/build-vms.nix` will generate NixOS configurations
-# for each node, and we need to include those configurations
-# in our Colmena setup as well.
+# TODO: Modularize most of this
 
 { insideVm ? false
 , deployers ? [ "deployer" ]           # Nodes configured as deployers (with Colmena and pre-built system closure)
@@ -26,13 +24,101 @@ let
 
   colmenaExec = "${colmena}/bin/colmena";
 
+  ## Utilities
   sshKeys = import (pkgs.path + "/nixos/tests/ssh-keys.nix") pkgs;
-  buildVms = import (pkgs.path + "/nixos/lib/build-vms.nix") {
-    inherit (pkgs) system pkgs lib;
+  nixosLib = import (pkgs.path + "/nixos/lib") { };
+
+  inputClosureOf = pkg: pkgs.runCommand "full-closure" {
+    refs = pkgs.writeReferencesToFile pkg.drvPath;
+  } ''
+    touch $out
+
+    while read ref; do
+      case $ref in
+        *.drv)
+          cat $ref >>$out
+          ;;
+      esac
+    done <$refs
+  '';
+
+  ## The modular NixOS test framework with Colmena additions
+  colmenaTestModule = { lib, config, ... }: let
+    cfg = config.colmena.test;
+
+    targetList = "[${concatStringsSep ", " targets}]";
+    bundle = pkgs.stdenv.mkDerivation {
+      name = "${config.name}-bundle";
+      dontUnpack = true;
+      dontInstall = true;
+      buildPhase = ''
+        cp -r ${cfg.bundle} $out
+        chmod u+w $out
+        cp ${./tools.nix} $out/tools.nix
+      '';
+    };
+  in {
+    options = {
+      colmena.test = {
+        bundle = lib.mkOption {
+          description = ''
+            Path to a directory to copy into the deployer as /tmp/bundle.
+          '';
+          type = lib.types.path;
         };
 
-  # Common setup
-  nodes = let
+        testScript = lib.mkOption {
+          description = ''
+            The test script.
+
+            The Colmena test framework will prepend initialization
+            statements to the actual test script.
+          '';
+          type = lib.types.str;
+        };
+      };
+    };
+    config = {
+      testScript = ''
+        start_all()
+      '' + lib.optionalString (prebuiltTarget != null) ''
+        deployer.succeed("nix-store -qR ${prebuiltSystem}")
+      '' + ''
+        deployer.succeed("nix-store -qR ${pkgs.path}")
+        deployer.succeed("ln -sf ${pkgs.path} /nixpkgs")
+        deployer.succeed("mkdir -p /root/.ssh && touch /root/.ssh/id_rsa && chmod 600 /root/.ssh/id_rsa && cat ${sshKeys.snakeOilPrivateKey} > /root/.ssh/id_rsa")
+
+        ${lib.optionalString (length targets != 0) ''
+        for node in ${targetList}:
+            node.wait_for_unit("sshd.service")
+            deployer.succeed(f"ssh -o StrictHostKeyChecking=accept-new {node.name} true", timeout=30)
+        ''}
+
+        deployer.succeed("cp --no-preserve=mode -r ${bundle} /tmp/bundle && chmod u+w /tmp/bundle")
+
+        orig_store_paths = set(deployer.succeed("ls /nix/store").strip().split("\n"))
+        def get_new_store_paths():
+            cur_store_paths = set(deployer.succeed("ls /nix/store").strip().split("\n"))
+            new_store_paths = cur_store_paths.difference(orig_store_paths)
+            deployer.log(f"{len(new_store_paths)} store paths were created")
+
+            l = list(map(lambda n: f"/nix/store/{n}", new_store_paths))
+            return l
+
+        ${cfg.testScript}
+      '';
+    };
+  };
+  evalTest = module: nixosLib.evalTest {
+    imports = [
+      module
+      colmenaTestModule
+      { hostPkgs = pkgs; }
+    ];
+  };
+
+  ## Common setup
+
   # Setup for deployer nodes
   #
   # We include the input closure of a prebuilt system profile
@@ -96,48 +182,34 @@ let
     virtualisation.writableStore = true;
   };
 
+  nodes = let
     deployerNodes = map (name: lib.nameValuePair name deployerConfig) deployers;
     targetNodes = map (name: lib.nameValuePair name targetConfig) targets;
   in listToAttrs (deployerNodes ++ targetNodes);
 
-  prebuiltSystem = let
-    all = buildVms.buildVirtualNetwork nodes;
-  in all.${prebuiltTarget}.config.system.build.toplevel;
+  # A "shallow" re-evaluation of the test for use from Colmena
+  standaloneTest = evalTest ({ ... }: {
+    inherit nodes;
+  });
 
-  # Utilities
-  getStandaloneConfigFor = node: let
-    configsWithIp = buildVms.assignIPAddresses nodes;
-  in { modulesPath, lib, config, ... }: {
-    imports = configsWithIp.${node} ++ [
-      (modulesPath + "/virtualisation/qemu-vm.nix")
-      (modulesPath + "/testing/test-instrumentation.nix")
+  prebuiltSystem = standaloneTest.config.nodes.${prebuiltTarget}.system.build.toplevel;
+
+  getStandaloneConfigFor = node: { lib, config, ... }: {
+    imports = [
+      (pkgs.path + "/nixos/lib/testing/nixos-test-base.nix")
+      (if elem node deployers then deployerConfig else targetConfig)
+      standaloneTest.config.nodes.${node}.system.build.networkConfig
     ];
 
     documentation.nixos.enable = lib.mkOverride 55 false;
     boot.loader.grub.enable = false;
     system.nixos.revision = lib.mkForce "constant-nixos-revision";
 
-    # otherwise the evaluation is unnecessarily slow in VM
-    virtualisation.additionalPaths = lib.mkForce [];
     nix.nixPath = lib.mkForce [ "nixpkgs=/nixpkgs" ];
 
     deployment.tags = lib.optional (config.networking.hostName != "deployer") "target";
   };
 
-  inputClosureOf = pkg: pkgs.runCommand "full-closure" {
-    refs = pkgs.writeReferencesToFile pkg.drvPath;
-  } ''
-    touch $out
-
-    while read ref; do
-      case $ref in
-        *.drv)
-          cat $ref >>$out
-          ;;
-      esac
-    done <$refs
-  '';
-
   makeTest = test: let
     customArgs = [ "bundle" ];
 
@@ -190,4 +262,9 @@ let
 in {
   inherit pkgs nodes colmena colmenaExec
     getStandaloneConfigFor inputClosureOf makeTest;
+
+  runTest = module: (evalTest ({ config, ... }: {
+    imports = [ module { inherit nodes; } ];
+    result = config.test;
+  })).config.result;
 }