From 372de58dffd0937b96731e17492c50f1016488b4 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Tue, 7 Dec 2021 23:13:31 -0800
Subject: [PATCH] integration-tests: Make it easier to override pkgs and
 colmena

---
 integration-tests/apply-local/default.nix |  4 +++-
 integration-tests/apply/default.nix       |  4 +++-
 integration-tests/apply/test-script.py    |  5 ++---
 integration-tests/default.nix             | 10 +++++-----
 integration-tests/exec/default.nix        |  4 +++-
 integration-tests/flakes/default.nix      |  4 +++-
 integration-tests/nixpkgs.nix             | 13 +++++++++++++
 integration-tests/parallel/default.nix    |  4 +++-
 integration-tests/tools.nix               | 12 +++---------
 9 files changed, 38 insertions(+), 22 deletions(-)
 create mode 100644 integration-tests/nixpkgs.nix

diff --git a/integration-tests/apply-local/default.nix b/integration-tests/apply-local/default.nix
index bd3346e..f3eeada 100644
--- a/integration-tests/apply-local/default.nix
+++ b/integration-tests/apply-local/default.nix
@@ -1,5 +1,7 @@
+{ pkgs ? import ../nixpkgs.nix }:
+
 let
-  tools = import ../tools.nix {
+  tools = pkgs.callPackage ../tools.nix {
     targets = [];
     prebuiltTarget = "deployer";
   };
diff --git a/integration-tests/apply/default.nix b/integration-tests/apply/default.nix
index 094e8f1..15b8d38 100644
--- a/integration-tests/apply/default.nix
+++ b/integration-tests/apply/default.nix
@@ -1,5 +1,7 @@
+{ pkgs ? import ../nixpkgs.nix }:
+
 let
-  tools = import ../tools.nix {};
+  tools = pkgs.callPackage ../tools.nix {};
 in tools.makeTest {
   name = "colmena-apply";
 
diff --git a/integration-tests/apply/test-script.py b/integration-tests/apply/test-script.py
index 73f994c..35daee9 100644
--- a/integration-tests/apply/test-script.py
+++ b/integration-tests/apply/test-script.py
@@ -22,9 +22,8 @@ with subtest("Check that activation messages were logged correctly"):
     assert "must appear during activation" in logs
 
 with subtest("Check that we can still connect to the target nodes"):
-    deployer.succeed("ssh alpha true")
-    deployer.succeed("ssh beta true")
-    deployer.succeed("ssh gamma true")
+    for node in targets:
+        deployer.succeed(f"ssh {node.name} true")
 
 with subtest("Check that the new configurations are indeed applied"):
     for node in targets:
diff --git a/integration-tests/default.nix b/integration-tests/default.nix
index b1b08a3..89b31c6 100644
--- a/integration-tests/default.nix
+++ b/integration-tests/default.nix
@@ -1,7 +1,7 @@
 {
-  apply = import ./apply;
-  apply-local = import ./apply-local;
-  exec = import ./exec;
-  flakes = import ./flakes;
-  parallel = import ./parallel;
+  apply = import ./apply {};
+  apply-local = import ./apply-local {};
+  exec = import ./exec {};
+  flakes = import ./flakes {};
+  parallel = import ./parallel {};
 }
diff --git a/integration-tests/exec/default.nix b/integration-tests/exec/default.nix
index 2ca1f6e..8020b23 100644
--- a/integration-tests/exec/default.nix
+++ b/integration-tests/exec/default.nix
@@ -1,5 +1,7 @@
+{ pkgs ? import ../nixpkgs.nix }:
+
 let
-  tools = import ../tools.nix {};
+  tools = pkgs.callPackage ../tools.nix {};
 in tools.makeTest {
   name = "colmena-exec";
 
diff --git a/integration-tests/flakes/default.nix b/integration-tests/flakes/default.nix
index d3865d6..c0ccc83 100644
--- a/integration-tests/flakes/default.nix
+++ b/integration-tests/flakes/default.nix
@@ -1,5 +1,7 @@
+{ pkgs ? import ../nixpkgs.nix }:
+
 let
-  tools = import ../tools.nix {
+  tools = pkgs.callPackage ../tools.nix {
     targets = [ "alpha" ];
   };
 in tools.makeTest {
diff --git a/integration-tests/nixpkgs.nix b/integration-tests/nixpkgs.nix
new file mode 100644
index 0000000..32781ba
--- /dev/null
+++ b/integration-tests/nixpkgs.nix
@@ -0,0 +1,13 @@
+let
+  lock = builtins.fromJSON (builtins.readFile ../flake.lock);
+  pinned = fetchTarball {
+    url = "https://github.com/NixOS/nixpkgs/archive/${lock.nodes.nixpkgs.locked.rev}.tar.gz";
+    sha256 = lock.nodes.nixpkgs.locked.narHash;
+  };
+in import pinned {
+  overlays = [
+    (final: prev: {
+      colmena = final.callPackage ../default.nix { };
+    })
+  ];
+}
diff --git a/integration-tests/parallel/default.nix b/integration-tests/parallel/default.nix
index 92eb14a..0d1a3c6 100644
--- a/integration-tests/parallel/default.nix
+++ b/integration-tests/parallel/default.nix
@@ -1,5 +1,7 @@
+{ pkgs ? import ../nixpkgs.nix }:
+
 let
-  tools = import ../tools.nix {};
+  tools = pkgs.callPackage ../tools.nix {};
 in tools.makeTest {
   name = "colmena-parallel";
 
diff --git a/integration-tests/tools.nix b/integration-tests/tools.nix
index 3f43e1b..19b1649 100644
--- a/integration-tests/tools.nix
+++ b/integration-tests/tools.nix
@@ -11,6 +11,9 @@
 , deployers ? [ "deployer" ]           # Nodes configured as deployers (with Colmena and pre-built system closure)
 , targets ? [ "alpha" "beta" "gamma" ] # Nodes configured as targets (minimal config)
 , prebuiltTarget ? "alpha"             # Target node to prebuild system closure for, or null
+
+, pkgs ? if insideVm then import <nixpkgs> {} else throw "Must specify pkgs"
+, colmena ? if !insideVm then pkgs.colmena else throw "Cannot eval inside VM"
 }:
 
 with builtins;
@@ -18,17 +21,8 @@ with builtins;
 assert elem "deployer" deployers;
 
 let
-  lock = builtins.fromJSON (builtins.readFile ../flake.lock);
-  pinned = if insideVm then <nixpkgs> else fetchTarball {
-    url = "https://github.com/NixOS/nixpkgs/archive/${lock.nodes.nixpkgs.locked.rev}.tar.gz";
-    sha256 = lock.nodes.nixpkgs.locked.narHash;
-  };
-  pkgs = import pinned {};
   inherit (pkgs) lib;
 
-  colmena =
-    if !insideVm then import ../default.nix { inherit pkgs; }
-    else throw "Cannot be used inside the VM";
   colmenaExec = "${colmena}/bin/colmena";
 
   sshKeys = import (pkgs.path + "/nixos/tests/ssh-keys.nix") pkgs;