nixos/tests: Implement python test-driver version of requireActiveUnit

1 files changed, 14 insertions, 1 deletions

diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py
index 45b7e229a5c..2c137ab3837 100644
--- a/nixos/lib/test-driver/test-driver.py
+++ b/nixos/lib/test-driver/test-driver.py
@@ -4,7 +4,9 @@ from contextlib import contextmanager
 from xml.sax.saxutils import XMLGenerator
 import _thread
 import atexit
+import json
 import os
+import ptpython.repl
 import pty
 import queue
 import re
@@ -15,7 +17,6 @@ import sys
 import tempfile
 import time
 import unicodedata
-import ptpython.repl
 
 CHAR_TO_KEY = {
     "A": "shift-a",
@@ -344,6 +345,18 @@ class Machine:
             )
         return self.execute("systemctl {}".format(q))
 
+    def require_unit_state(self, unit, require_state="active"):
+        with self.nested(
+            "checking if unit ‘{}’ has reached state '{}'".format(unit, require_state)
+        ):
+            info = self.get_unit_info(unit)
+            state = info["ActiveState"]
+            if state != require_state:
+                raise Exception(
+                    "Expected unit ‘{}’ to to be in state ".format(unit)
+                    + "'active' but it is in state ‘{}’".format(state)
+                )
+
     def execute(self, command):
         self.connect()