add ubuntu kexec test

Author: Mic92

flake.lock

@@ -11,6 +11,8 @@
     nixos-images.url = "github:nix-community/nixos-images";
     nixos-images.inputs.nixos-unstable.follows = "nixpkgs";
     nixos-images.inputs.nixos-stable.follows = "nixos-stable";
+    # https://github.com/numtide/nix-vm-test/pull/105
+    nix-vm-test = { url = "github:Mic92/nix-vm-test"; inputs.nixpkgs.follows = "nixpkgs"; };
 
     # used for development
     treefmt-nix = { url = "github:numtide/treefmt-nix"; inputs.nixpkgs.follows = "nixpkgs"; };

flake.nix

@@ -17,6 +17,9 @@
       testInputsStable = testInputsUnstable // {
         kexec-installer = "${inputs'.nixos-images.packages.kexec-installer-nixos-stable-noninteractive}/nixos-kexec-installer-noninteractive-${system}.tar.gz";
       };
+      ubuntuTestInputs = testInputsUnstable // {
+        nix-vm-test = inputs.nix-vm-test;
+      };
     in
     {
       from-nixos = import ./from-nixos.nix testInputsUnstable;
@@ -26,5 +29,6 @@
       from-nixos-with-generated-config = import ./from-nixos-generate-config.nix testInputsUnstable;
       from-nixos-build-on-remote = import ./from-nixos-build-on-remote.nix testInputsUnstable;
       from-nixos-separated-phases = import ./from-nixos-separated-phases.nix testInputsUnstable;
+      ubuntu-kexec-test = import ./ubuntu-kexec-test.nix ubuntuTestInputs;
     });
 }

tests/flake-module.nix

@@ -1,76 +1,97 @@
-{ pkgs, nixos-anywhere, kexec-installer, nix-vm-test, ... }:
-
-let
-  # Create dummy store paths for testing
-  testPaths = pkgs.runCommand "test-store-paths" {} ''
-    mkdir -p $out
-    echo "echo 'Dummy disko script'" > $out/disko
-    echo "echo 'Dummy system closure'" > $out/system
-    chmod +x $out/disko $out/system
-  '';
-in
-nix-vm-test.ubuntu."22_04" {
-  memorySize = 2048;
-
-  # Forward SSH port to allow connection from the host
-  virtualisation.forwardPorts = [
-    { from = "host"; host.port = 2222; guest.port = 22; }
-  ];
-
-  # Make the SSH keys available in the VM
-  sharedDirs = {
-    sshKeys = {
-      source = "${nixos-anywhere}/tests/modules/ssh-keys";
-      target = "/tmp/ssh-keys";
-    };
+{ pkgs, nixos-anywhere, kexec-installer, nix-vm-test, system-to-install, ... }:
+
+(nix-vm-test.lib.${pkgs.system}.ubuntu."24_04" {
+  sharedDirs = { };
+
+  # Configure VM with 2GB memory
+  machineConfigModule = { ... }: {
+    nodes.vm.virtualisation.memorySize = 2048;
   };
 
   # The test script
   testScript = ''
+    # Python imports
+    import subprocess
+    import tempfile
+    import shutil
+    import os
+
     # Wait for the system to be fully booted
     vm.wait_for_unit("multi-user.target")
 
     # Unmask SSH service (which is masked by default in the test VM)
-    vm.succeed("systemctl unmask ssh.service")
-    vm.succeed("systemctl unmask ssh.socket")
-    vm.succeed("systemctl start ssh")
+    vm.succeed("systemctl unmask ssh.service ssh.socket")
+
+    # Generate SSH host keys (required for SSH to start)
+    vm.succeed("ssh-keygen -A")
 
     # Setup SSH with the existing keys
     vm.succeed("mkdir -p /root/.ssh")
-    vm.succeed("cp /tmp/ssh-keys/ssh.pub /root/.ssh/authorized_keys")
+    vm.succeed(
+      "echo '${builtins.replaceStrings ["\n"] [""] (builtins.readFile ./modules/ssh-keys/ssh.pub)}' > /root/.ssh/authorized_keys"
+    )
     vm.succeed("chmod 644 /root/.ssh/authorized_keys")
 
     # Setup SSH for connection from host
-    vm.succeed("sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config")
-    vm.succeed("systemctl restart ssh")
-    
+    vm.succeed(
+      "sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config"
+    )
+
+    # Start SSH service
+    vm.succeed("systemctl start ssh")
+
     # Wait for SSH to be available
     vm.wait_for_open_port(22)
-    print("SSH server is ready")
 
-    # Run nixos-anywhere from the host against the VM
-    print("Starting kexec test phase...")
-    
-    # Use Python's subprocess to run nixos-anywhere from the host
-    import subprocess
-    
-    cmd = f"""${nixos-anywhere}/bin/nixos-anywhere \\
-      -i ${nixos-anywhere}/tests/modules/ssh-keys/ssh \\
-      --ssh-port 2222 \\
-      --phases kexec \\
-      --kexec ${kexec-installer} \\
-      --store-paths ${testPaths}/disko ${testPaths}/system \\
-      --debug \\
-      root@localhost
-    """
-    
-    print(f"Running command: {cmd}")
-    result = subprocess.run(cmd, shell=True, check=False)
-    
-    if result.returncode == 0:
-      print("kexec phase completed successfully")
-    else:
-      print(f"nixos-anywhere failed with exit code {result.returncode}")
-      raise Exception("nixos-anywhere command failed")
+    # Forward SSH port using vm.forward_port method
+    ssh_port = 2222
+    vm.forward_port(host_port=ssh_port, guest_port=22)
+
+    # Use temporary file for SSH key with automatic cleanup
+    with tempfile.NamedTemporaryFile(mode='w', delete=True, suffix='_ssh_key') as temp_key:
+      temp_key_path = temp_key.name
+
+      # Copy SSH private key to temp file with correct permissions
+      shutil.copy2("${./modules/ssh-keys/ssh}", temp_key_path)
+      os.chmod(temp_key_path, 0o600)
+
+      nixos_anywhere_cmd = [
+        "${nixos-anywhere}/bin/nixos-anywhere",
+        "-i", temp_key_path,
+        "--ssh-port", str(ssh_port),
+        "--post-kexec-ssh-port", "2222",
+        "--phases", "kexec",
+        "--kexec", "${kexec-installer}",
+        "--store-paths", "${system-to-install.config.system.build.diskoScriptNoDeps}",
+        "${system-to-install.config.system.build.toplevel}",
+        "--debug",
+        "root@localhost"
+      ]
+
+      result = subprocess.run(nixos_anywhere_cmd, check=False)
+
+      if result.returncode != 0:
+        print(f"nixos-anywhere failed with exit code {result.returncode}")
+        vm.succeed("dmesg | tail -n 50")
+        vm.succeed("journalctl -n 50")
+        raise Exception(f"nixos-anywhere command failed with exit code {result.returncode}")
+
+      # Test SSH connection to verify we're in NixOS kexec environment
+      check_cmd = [
+        "${pkgs.openssh}/bin/ssh", "-v",
+        "-i", temp_key_path,
+        "-p", "2222",
+        "-o", "StrictHostKeyChecking=no",
+        "-o", "UserKnownHostsFile=/dev/null",
+        "root@localhost",
+        "cat /etc/os-release"
+      ]
+
+      test_result = subprocess.run(check_cmd, check=True, stdout=subprocess.PIPE, text=True)
+      assert "nixos" in test_result.stdout.lower(), f"Expected NixOS environment but got: {test_result.stdout}"
+
+    # After kexec we no longer have the machine driver,
+    # so we need to let the VM crash because the test driver backdoor gets confused by the terminal output.
+    vm.crash()
   '';
-}
+}).sandboxed