diff --git a/nixos/lib/test-driver/src/test_driver/machine/__init__.py b/nixos/lib/test-driver/src/test_driver/machine/__init__.py
index ffc7c7517165..596627a79491 100644
--- a/nixos/lib/test-driver/src/test_driver/machine/__init__.py
+++ b/nixos/lib/test-driver/src/test_driver/machine/__init__.py
@@ -335,7 +335,7 @@ class BaseMachine(ABC):
         ...
 
     @abstractmethod
-    def wait_for_shutdown(self) -> None:
+    def wait_for_shutdown(self, timeout: dt.timedelta | None = None) -> None:
         """Wait for the machine to power off. This does *not* initiate a shutdown;
         that's usually done via `shutdown()`.
         """
@@ -1061,7 +1061,7 @@ class QemuMachine(BaseMachine):
                 break
             self.send_console(char.decode())
 
-    def wait_for_shutdown(self) -> None:
+    def wait_for_shutdown(self, timeout: dt.timedelta | None = None) -> None:
         """
         Wait for the VM to power off. This does *not* initiate a shutdown;
         that's usually done via `shutdown()`.
@@ -1072,7 +1072,9 @@ class QemuMachine(BaseMachine):
         with self.nested("waiting for the VM to power off"):
             sys.stdout.flush()
             assert self.process
-            self.process.wait()
+            self.process.wait(
+                timeout=timeout.total_seconds() if timeout is not None else None
+            )
 
             self.pid = None
             self.booted = False
@@ -1903,7 +1905,7 @@ class NspawnMachine(BaseMachine):
         self.systemctl("poweroff")
         self.wait_for_shutdown()
 
-    def wait_for_shutdown(self) -> None:
+    def wait_for_shutdown(self, timeout: dt.timedelta | None = None) -> None:
         """
         Wait for the container to power off. This does *not* initiate a shutdown;
         that's usually done via `shutdown()`.
@@ -1912,7 +1914,9 @@ class NspawnMachine(BaseMachine):
             return
 
         with self.nested("waiting for the container to power off"):
-            self.process.wait()
+            self.process.wait(
+                timeout=timeout.total_seconds() if timeout is not None else None
+            )
             self.process = None
 
 
diff --git a/nixos/tests/systemd-shutdown.nix b/nixos/tests/systemd-shutdown.nix
index 7439c9331182..76870b2e8562 100644
--- a/nixos/tests/systemd-shutdown.nix
+++ b/nixos/tests/systemd-shutdown.nix
@@ -13,24 +13,32 @@ in
 
   nodes.machine = {
     imports = [ ../modules/profiles/minimal.nix ];
+    systemd.shutdown.pre-exitrd = pkgs.writeShellScript "pre-exitrd" ''
+      echo pre-exitrd > /run/initramfs/test.txt
+    '';
     systemd.shutdownRamfs.contents."/etc/systemd/system-shutdown/shutdown-message".source =
       pkgs.writeShellScript "shutdown-message" ''
-        echo "${msg}" > /dev/kmsg
+        if test -e /test.txt; then
+          # Test should only pass if both scripts run.
+          echo "${msg}" > /dev/kmsg
+        fi
       '';
     boot.initrd.systemd.enable = systemdStage1;
   };
 
   testScript = ''
+    import datetime as dt
+
     # Check that 'generate-shutdown-ramfs.service' is started
     # automatically and that 'systemd-shutdown' runs our script.
     machine.wait_for_unit("multi-user.target")
     # .shutdown() would wait for the machine to power off
     machine.execute("systemctl poweroff", check_return=False)
     # Message printed by systemd-shutdown
-    machine.wait_for_console_text("Unmounting '/oldroot'")
-    machine.wait_for_console_text("${msg}")
+    machine.wait_for_console_text("Unmounting '/oldroot'", timeout=dt.timedelta(seconds=60))
+    machine.wait_for_console_text("${msg}", timeout=dt.timedelta(seconds=5))
     # Don't try to sync filesystems
-    machine.wait_for_shutdown()
+    machine.wait_for_shutdown(timeout=dt.timedelta(seconds=5))
 
     # In a separate boot, start 'generate-shutdown-ramfs.service'
     # manually in order to check the permissions on '/run/initramfs'.