test-driver: doctest some pure functions

nixos/lib/test-driver/src/test_driver/driver.py

@@ -87,6 +87,30 @@ def get_tmp_dir() -> Path:
 
 
 def pythonize_name(name: str) -> str:
+    """Turn a machine name into a valid Python identifier.
+
+    Names that already are valid identifiers stay the same:
+
+    >>> pythonize_name("node_1")
+    'node_1'
+
+    Every character that is not allowed in an identifier becomes "_":
+
+    >>> pythonize_name("my-vm")
+    'my_vm'
+    >>> pythonize_name("web.example.org")
+    'web_example_org'
+    >>> pythonize_name("münchen")
+    'm_nchen'
+
+    A leading digit is replaced rather than prefixed, so it gets lost, while
+    digits later in the name are kept:
+
+    >>> pythonize_name("1node")
+    '_node'
+    >>> pythonize_name("12node")
+    '_2node'
+    """
     return re.sub(r"^[^A-Za-z_]|[^A-Za-z0-9_]", "_", name)
 
 

nixos/lib/test-driver/src/test_driver/duration.py

@@ -14,6 +14,22 @@ def as_timedelta(duration: Duration) -> dt.timedelta:
     Bare numbers are interpreted as seconds. This keeps existing test scripts
     that pass plain integers working while `datetime.timedelta` becomes the preferred
     way to express durations.
+
+    >>> as_timedelta(dt.timedelta(minutes=1))
+    datetime.timedelta(seconds=60)
+    >>> as_timedelta(90)
+    datetime.timedelta(seconds=90)
+    >>> as_timedelta(1.5)
+    datetime.timedelta(seconds=1, microseconds=500000)
+
+    A `bool` is an `int` in Python, but it is never a valid duration:
+
+    >>> as_timedelta(True)
+    Traceback (most recent call last):
+    TypeError: expected a duration, got bool: True
+    >>> as_timedelta("1")
+    Traceback (most recent call last):
+    TypeError: expected a timedelta, int, or float duration, got str
     """
     if isinstance(duration, dt.timedelta):
         return duration
@@ -27,7 +43,13 @@ def as_timedelta(duration: Duration) -> dt.timedelta:
 
 
 def as_seconds(duration: Duration) -> float:
-    """Coerce a `Duration` into a floating-point number of seconds."""
+    """Coerce a `Duration` into a floating-point number of seconds.
+
+    >>> as_seconds(90)
+    90.0
+    >>> as_seconds(dt.timedelta(minutes=1, milliseconds=500))
+    60.5
+    """
     return as_timedelta(duration).total_seconds()
 
 

nixos/lib/test-driver/src/test_driver/logger.py

@@ -322,10 +322,27 @@ class XMLLogger(AbstractLogger):
         self.xml.endDocument()
         self.logfile_handle.close()
 
-    def sanitise(self, message: str) -> str:
+    @staticmethod
+    def sanitise(message: str) -> str:
+        r"""Remove the characters of Unicode's "Other" categories, such as
+        control characters, from serial output before it is written to XML.
+
+        >>> XMLLogger.sanitise("booting\x00...\x07")
+        'booting...'
+        >>> XMLLogger.sanitise("Grüße")
+        'Grüße'
+        """
         return "".join(ch for ch in message if unicodedata.category(ch)[0] != "C")
 
-    def maybe_prefix(self, message: str, attributes: dict[str, str]) -> str:
+    @staticmethod
+    def maybe_prefix(message: str, attributes: dict[str, str]) -> str:
+        """Prefix the message with the name of the machine, if there is one.
+
+        >>> XMLLogger.maybe_prefix("booting", {"machine": "server"})
+        'server: booting'
+        >>> XMLLogger.maybe_prefix("booting", {})
+        'booting'
+        """
         if "machine" in attributes:
             return f"{attributes['machine']}: {message}"
         return message

nixos/lib/test-driver/src/test_driver/machine/__init__.py

@@ -139,7 +139,22 @@ XDOTOOL_MODIFIER_ALIASES = {
 def xdotool_key_sequence(key: str) -> str:
     """Translate a QEMU-style key chord to xdotool syntax.
 
-    For example, "ctrl-alt-ret" becomes "ctrl+alt+Return".
+    >>> xdotool_key_sequence("ctrl-alt-ret")
+    'ctrl+alt+Return'
+    >>> xdotool_key_sequence("meta_l-x")
+    'super+x'
+    >>> xdotool_key_sequence("alt-f4")
+    'alt+F4'
+
+    A modifier on its own is a regular key:
+
+    >>> xdotool_key_sequence("shift")
+    'Shift_L'
+
+    A "-" that follows a modifier is the key itself, not a separator:
+
+    >>> xdotool_key_sequence("ctrl--")
+    'ctrl+-'
     """
     parts = key.split("-")
     modifiers = []
@@ -154,6 +169,13 @@ def xdotool_key_sequence(key: str) -> str:
 
 
 def make_command(args: list) -> str:
+    """Join arguments into a shell command, quoting each one as a single word.
+
+    >>> make_command(["echo", "a b", 1])
+    "echo 'a b' 1"
+    >>> make_command(["ls", "*"])
+    "ls '*'"
+    """
     return " ".join(map(shlex.quote, (map(str, args))))
 
 
@@ -217,6 +239,13 @@ class QemuStartCommand:
         These Nix commands have the particular characteristic that the
         machine name can be extracted out of them via a regex match.
         (Admittedly a _very_ implicit contract, TODO fix this eventually.)
+
+        >>> QemuStartCommand("/nix/store/hash-nixos-vm/bin/run-foo-vm").machine_name
+        'foo'
+        >>> QemuStartCommand("run-foo-bar-vm").machine_name
+        'foo-bar'
+        >>> QemuStartCommand("/some/script").machine_name
+        'machine'
         """
         match = re.search("run-(.+)-vm$", self._cmd)
         name = "machine"
@@ -1690,7 +1719,19 @@ class NspawnMachine(BaseMachine):
     vnc: NspawnVnc | None
 
     @staticmethod
-    def machine_name_from_start_command(start_command: str) -> str:
+    def _machine_name_from_start_command(start_command: str) -> str:
+        """Extract the machine name from a start script of a container.
+
+        >>> NspawnMachine._machine_name_from_start_command(
+        ...     "/nix/store/hash-container/bin/run-foo-nspawn"
+        ... )
+        'foo'
+        >>> NspawnMachine._machine_name_from_start_command("run-foo-bar-nspawn")
+        'foo-bar'
+        >>> NspawnMachine._machine_name_from_start_command("/some/script")
+        Traceback (most recent call last):
+        AssertionError: Could not extract node name from /some/script
+        """
         match = re.search("run-(.+)-nspawn", os.path.basename(start_command))
         assert match is not None, f"Could not extract node name from {start_command}"
         return match.group(1)
@@ -1710,7 +1751,7 @@ class NspawnMachine(BaseMachine):
         # See analogous TODO in `QemuStartCommand::machine_name`.
         super().__init__(
             out_dir=out_dir,
-            name=name or self.machine_name_from_start_command(start_command),
+            name=name or self._machine_name_from_start_command(start_command),
             logger=logger,
             callbacks=callbacks,
             tmp_dir=tmp_dir,

nixos/lib/test-driver/src/test_driver/machine/qmp.py

@@ -12,6 +12,23 @@ logger = logging.getLogger(__name__)
 
 
 class QMPAPIError(RuntimeError):
+    """An error response from the QMP API.
+
+    >>> error = {"class": "GenericError", "desc": "boom"}
+    >>> print(QMPAPIError({"error": error, "id": 3}))
+    <QMP API error related to transaction 3 [GenericError]: boom>
+
+    The transaction id is absent from errors that the server hits before it
+    could read the id:
+
+    >>> print(QMPAPIError({"error": error}))
+    <QMP API error related to transaction None [GenericError]: boom>
+
+    >>> QMPAPIError({"error": {"class": "GenericError"}})
+    Traceback (most recent call last):
+    RuntimeError: Malformed QMP API error response
+    """
+
     def __init__(self, message: dict[str, Any]):
         assert "error" in message, "Not an error message!"
         error = message["error"]

nixos/lib/test-driver/src/test_driver/polling_condition.py

@@ -13,6 +13,35 @@ class PollingConditionError(Exception):
 
 
 class PollingCondition:
+    """A condition that gets checked periodically while a test runs.
+
+    The description comes from the `description` argument. Without one, it
+    falls back to the docstring of the condition, and then to its name:
+
+    >>> from test_driver.logger import CompositeLogger
+    >>> logger = CompositeLogger([])
+    >>> def service_runs():
+    ...     "The service is running"
+    >>> def no_docstring():
+    ...     pass
+    >>> PollingCondition(service_runs, logger, description="custom").description
+    'custom'
+    >>> PollingCondition(service_runs, logger).description
+    'The service is running'
+    >>> PollingCondition(no_docstring, logger).description
+    'no_docstring'
+
+    Log messages include the description:
+
+    >>> condition = PollingCondition(service_runs, logger)
+    >>> condition.nested_message
+    "Checking polling condition 'The service is running'"
+    >>> condition.status_message(True)
+    'Polling condition succeeded: The service is running'
+    >>> condition.status_message(False)
+    'Polling condition failed: The service is running'
+    """
+
     condition: Callable[[], bool]
     interval: dt.timedelta
     description: str | None