NixOS test framework: add overriding methods (#421480)

nixos/lib/testing/call-test.nix

@@ -1,6 +1,19 @@
-{ config, lib, ... }:
+{
+  config,
+  extendModules,
+  lib,
+  ...
+}:
 let
   inherit (lib) mkOption types;
+
+  unsafeGetAttrPosStringOr =
+    default: name: value:
+    let
+      p = builtins.unsafeGetAttrPos name value;
+    in
+    if p == null then default else p.file + ":" + toString p.line + ":" + toString p.column;
+
 in
 {
   options = {
@@ -9,4 +22,21 @@ in
       default = config;
     };
   };
+  config = {
+    # Docs: nixos/doc/manual/development/writing-nixos-tests.section.md
+    /**
+      See https://nixos.org/manual/nixos/unstable#sec-override-nixos-test
+    */
+    passthru.extend =
+      args@{
+        modules,
+        specialArgs ? { },
+      }:
+      (extendModules {
+        inherit specialArgs;
+        modules = map (lib.setDefaultModuleLocation (
+          unsafeGetAttrPosStringOr "<test.extend module>" "modules" args
+        )) modules;
+      }).config.test;
+  };
 }

nixos/lib/testing/nodes.nix

@@ -3,6 +3,7 @@ testModuleArgs@{
   lib,
   hostPkgs,
   nodes,
+  options,
   ...
 }:
 
@@ -73,6 +74,9 @@ let
     ];
   };
 
+  # TODO (lib): Dedup with run.nix, add to lib/options.nix
+  mkOneUp = opt: f: lib.mkOverride (opt.highestPrio - 1) (f opt.value);
+
 in
 
 {
@@ -233,5 +237,23 @@ in
       ))
     ];
 
+    # Docs: nixos/doc/manual/development/writing-nixos-tests.section.md
+    /**
+      See https://nixos.org/manual/nixos/unstable#sec-override-nixos-test
+    */
+    passthru.extendNixOS =
+      {
+        module,
+        specialArgs ? { },
+      }:
+      config.passthru.extend {
+        modules = [
+          {
+            extraBaseModules = module;
+            node.specialArgs = mkOneUp options.node.specialArgs (_: specialArgs);
+          }
+        ];
+      };
+
   };
 }

nixos/lib/testing/run.nix

@@ -2,10 +2,31 @@
   config,
   hostPkgs,
   lib,
+  options,
   ...
 }:
 let
   inherit (lib) types mkOption;
+
+  # TODO (lib): Also use lib equivalent in nodes.nix
+  /**
+    Create a module system definition that overrides an existing option from a different module evaluation.
+
+    Type: Option a -> (a -> a) -> Definition a
+  */
+  mkOneUp =
+    /**
+      Option from an existing module evaluation, e.g.
+      - `(lib.evalModules ...).options.x` when invoking `evalModules` again,
+      - or `{ options, ... }:` when invoking `extendModules`.
+    */
+    opt:
+    /**
+      Function from the old value to the new definition, which will be wrapped with `mkOverride`.
+    */
+    f:
+    lib.mkOverride (opt.highestPrio - 1) (f opt.value);
+
 in
 {
   options = {
@@ -30,6 +51,13 @@ in
       internal = true;
     };
 
+    rawTestDerivationArg = mkOption {
+      type = types.functionTo types.raw;
+      description = ''
+        Argument passed to `mkDerivation` to create the `rawTestDerivation`.
+      '';
+    };
+
     test = mkOption {
       type = types.package;
       # TODO: can the interactive driver be configured to access the network?
@@ -43,10 +71,12 @@ in
   };
 
   config = {
-    rawTestDerivation =
+    rawTestDerivation = hostPkgs.stdenv.mkDerivation config.rawTestDerivationArg;
+    rawTestDerivationArg =
+      finalAttrs:
       assert lib.assertMsg (!config.sshBackdoor.enable)
         "The SSH backdoor is currently not supported for non-interactive testing! Please make sure to only set `interactive.sshBackdoor.enable = true;`!";
-      hostPkgs.stdenv.mkDerivation {
+      {
         name = "vm-test-run-${config.name}";
 
         requiredSystemFeatures =
@@ -75,5 +105,19 @@ in
 
     # useful for inspection (debugging / exploration)
     passthru.config = config;
+
+    # Docs: nixos/doc/manual/development/writing-nixos-tests.section.md
+    /**
+      See https://nixos.org/manual/nixos/unstable#sec-override-nixos-test
+    */
+    passthru.overrideTestDerivation =
+      f:
+      config.passthru.extend {
+        modules = [
+          {
+            rawTestDerivationArg = mkOneUp options.rawTestDerivationArg (lib.extends (lib.toExtension f));
+          }
+        ];
+      };
   };
 }