From 9d8f0f5e5c8e79fb51fa3299664a2401517a717a Mon Sep 17 00:00:00 2001
From: Nathan Henrie <nate@n8henrie.com>
Date: Tue, 21 Feb 2023 12:18:25 -0700
Subject: [PATCH] Migrate testing into python module

This allows using python linters and tools directly and makes it much
easier to create our own convenience methods.

If we can avoid the trap of excess complexity and configurability, I
think this could facilitate addition of much broader integration tests,
especially for the CLI.
---
 test/agenix_testing.py | 130 +++++++++++++++++++++++++++++++++++++++++
 test/integration.nix   |  68 +++++++++------------
 test/pyproject.toml    |   7 +++
 3 files changed, 165 insertions(+), 40 deletions(-)
 create mode 100644 test/agenix_testing.py
 create mode 100644 test/pyproject.toml

diff --git a/test/agenix_testing.py b/test/agenix_testing.py
new file mode 100644
index 0000000..c672a0b
--- /dev/null
+++ b/test/agenix_testing.py
@@ -0,0 +1,130 @@
+"""Provide a class and helper methods for agenix integration tests."""
+
+
+import typing as t
+
+T = t.TypeVar("T", str, list[str])
+
+
+class AgenixTester:
+    """Provide a class to help reduce repetition in setup."""
+
+    def __init__(self, system, user, password) -> None:
+        """Necessary setup can be put here."""
+        self.system = system
+        self.user = user
+        self.password = password
+        self.setup()
+
+    def login(self) -> None:
+        self.system.wait_for_unit("multi-user.target")
+        self.system.wait_until_succeeds("pgrep -f 'agetty.*tty1'")
+        self.system.sleep(2)
+        self.system.send_key("alt-f2")
+        self.system.wait_until_succeeds("[ $(fgconsole) = 2 ]")
+        self.system.wait_for_unit("getty@tty2.service")
+        self.system.wait_until_succeeds("pgrep -f 'agetty.*tty2'")
+        self.system.wait_until_tty_matches("2", "login: ")
+        self.system.send_chars(f"{self.user}\n")
+        self.system.wait_until_tty_matches("2", f"login: {self.user}")
+        self.system.wait_until_succeeds("pgrep login")
+        self.system.sleep(2)
+        self.system.send_chars(f"{self.password}\n")
+
+    def setup(self) -> None:
+        """Run common setup code."""
+        self.login()
+
+    def user_succeed(
+        self,
+        cmds: T,
+        directory: str | None = None,
+        debug: bool = False,
+    ) -> T:
+        """Run cmds as `self.user`, optionally in a specified directory.
+
+        For convenience, if cmds is a sequence, returns output as a list of
+        outputs corresponding with each line in cmds. if cmds is a string,
+        returns output as a string.
+        """
+
+        context: list[str] = [
+            "set -Eeu -o pipefail",
+            "shopt -s inherit_errexit",
+        ]
+        if debug:
+            context.append("set -x")
+
+        if directory:
+            context.append(f"cd {directory}")
+
+        if isinstance(cmds, str):
+            commands_str = "\n".join([*context, cmds])
+            final_command = f"sudo -u {self.user} -- bash -c '{commands_str}'"
+            return self.system.succeed(final_command)
+
+        results: list[str] = []
+        for cmd in cmds:
+            commands_str = "\n".join([*context, cmd])
+            final_command = f"sudo -u {self.user} -- bash -c '{commands_str}'"
+            result = self.system.succeed(final_command)
+            results.append(result.strip())
+        return t.cast(T, results)
+
+    def run_all(self) -> None:
+        self.test_rekeying()
+        self.test_user_edit()
+
+    def test_rekeying(self) -> None:
+        """Ensure we can rekey a file and its hash changes."""
+
+        before_hash, _, after_hash = self.user_succeed(
+            [
+                "sha256sum passwordfile-user1.age",
+                f"agenix -r -i /home/{self.user}/.ssh/id_ed25519",
+                "sha256sum passwordfile-user1.age",
+            ],
+            directory="/tmp/secrets",
+        )
+
+        # Ensure we actually have hashes
+        for line in [before_hash, after_hash]:
+            h = line.split()
+            assert len(h) == 2, f"hash should be [hash, filename], got {h}"
+            assert h[1] == "passwordfile-user1.age", "filename is incorrect"
+            assert len(h[0].strip()) == 64, "hash length is incorrect"
+        assert (
+            before_hash[0] != after_hash[0]
+        ), "hash did not change with rekeying"
+
+    def test_user_edit(self):
+        """Ensure user1 can edit passwordfile-user1.age."""
+        self.user_succeed(
+            "EDITOR=cat agenix -e passwordfile-user1.age",
+            directory="/tmp/secrets",
+        )
+
+        self.user_succeed("echo bogus > ~/.ssh/id_rsa")
+
+        # Cannot edit with bogus default id_rsa
+        self.system.fail(
+            f"sudo -u {self.user} -- bash -c '"
+            "cd /tmp/secrets; "
+            "EDITOR=cat agenix -e /tmp/secrets/passwordfile-user1.age; "
+            "'"
+        )
+
+        # user1 can still edit if good identity specified
+        *_, pw = self.user_succeed(
+            [
+                (
+                    "EDITOR=cat agenix -e passwordfile-user1.age "
+                    "-i /home/user1/.ssh/id_ed25519"
+                ),
+                "rm ~/.ssh/id_rsa",
+                "echo 'secret1234' | agenix -e passwordfile-user1.age",
+                "EDITOR=cat agenix -e passwordfile-user1.age",
+            ],
+            directory="/tmp/secrets",
+        )
+        assert pw == "secret1234", f"password didn't match, got '{pw}'"
diff --git a/test/integration.nix b/test/integration.nix
index ff1bbac..e90ad97 100644
--- a/test/integration.nix
+++ b/test/integration.nix
@@ -9,6 +9,24 @@
 }:
 pkgs.nixosTest {
   name = "agenix-integration";
+  extraPythonPackages = ps: let
+    agenixTesting = let
+      version = (pkgs.callPackage ../pkgs/agenix.nix {}).version;
+    in
+      ps.buildPythonPackage rec {
+        inherit version;
+        pname = "agenix_testing";
+        src = ./.;
+        format = "pyproject";
+        propagatedBuildInputs = [ps.setuptools];
+        postPatch = ''
+          # Keep a default version makes for easy installation outside of
+          # nix for debugging
+          substituteInPlace pyproject.toml \
+            --replace 'version = "0.1.0"' 'version = "${version}"'
+        '';
+      };
+  in [agenixTesting];
   nodes.system1 = {
     config,
     pkgs,
@@ -49,47 +67,17 @@ pkgs.nixosTest {
     user = "user1";
     password = "password1234";
   in ''
-    system1.wait_for_unit("multi-user.target")
-    system1.wait_until_succeeds("pgrep -f 'agetty.*tty1'")
-    system1.sleep(2)
-    system1.send_key("alt-f2")
-    system1.wait_until_succeeds("[ $(fgconsole) = 2 ]")
-    system1.wait_for_unit("getty@tty2.service")
-    system1.wait_until_succeeds("pgrep -f 'agetty.*tty2'")
-    system1.wait_until_tty_matches("2", "login: ")
-    system1.send_chars("${user}\n")
-    system1.wait_until_tty_matches("2", "login: ${user}")
-    system1.wait_until_succeeds("pgrep login")
-    system1.sleep(2)
-    system1.send_chars("${password}\n")
+    # Skipping analyzing "agenix_testing": module is installed, but missing
+    # library stubs or py.typed marker
+    from agenix_testing import AgenixTester  # type: ignore
+    tester = AgenixTester(system=system1, user="${user}", password="${password}")
+
+    # Can still be used as before
     system1.send_chars("whoami > /tmp/1\n")
-    system1.wait_for_file("/tmp/1")
-    assert "${user}" in system1.succeed("cat /tmp/1")
+    # Or from `tester.system`
+    tester.system.wait_for_file("/tmp/1")
+    assert "${user}" in tester.system.succeed("cat /tmp/1")
 
-    userDo = lambda input : f"sudo -u user1 -- bash -c 'set -eou pipefail; cd /tmp/secrets; {input}'"
-
-    before_hash = system1.succeed(userDo('sha256sum passwordfile-user1.age')).split()
-    print(system1.succeed(userDo('agenix -r -i /home/user1/.ssh/id_ed25519')))
-    after_hash = system1.succeed(userDo('sha256sum passwordfile-user1.age')).split()
-
-    # Ensure we actually have hashes
-    for h in [before_hash, after_hash]:
-        assert len(h) == 2, "hash should be [hash, filename]"
-        assert h[1] == "passwordfile-user1.age", "filename is incorrect"
-        assert len(h[0].strip()) == 64, "hash length is incorrect"
-    assert before_hash[0] != after_hash[0], "hash did not change with rekeying"
-
-    # user1 can edit passwordfile-user1.age
-    system1.succeed(userDo("EDITOR=cat agenix -e passwordfile-user1.age"))
-
-    # user1 can edit even if bogus id_rsa present
-    system1.succeed(userDo("echo bogus > ~/.ssh/id_rsa"))
-    system1.fail(userDo("EDITOR=cat agenix -e passwordfile-user1.age"))
-    system1.succeed(userDo("EDITOR=cat agenix -e passwordfile-user1.age -i /home/user1/.ssh/id_ed25519"))
-    system1.succeed(userDo("rm ~/.ssh/id_rsa"))
-
-    # user1 can edit a secret by piping in contents
-    system1.succeed(userDo("echo 'secret1234' | agenix -e passwordfile-user1.age"))
-    assert "secret1234" in system1.succeed(userDo("EDITOR=cat agenix -e passwordfile-user1.age"))
+    tester.run_all()
   '';
 }
diff --git a/test/pyproject.toml b/test/pyproject.toml
new file mode 100644
index 0000000..55e0890
--- /dev/null
+++ b/test/pyproject.toml
@@ -0,0 +1,7 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "agenix_testing"
+version = "0.1.0"