nixos/test: colour machine names

1 files changed, 13 insertions, 2 deletions

diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py
index 99c6eb06de3c8..dc11eaa1982ff 100644
--- a/nixos/lib/test-driver/test-driver.py
+++ b/nixos/lib/test-driver/test-driver.py
@@ -3,6 +3,7 @@ import argparse
 import atexit
 import base64
 import io
+import itertools
 import logging
 import os
 import pathlib
@@ -92,10 +93,17 @@ logging.basicConfig(format="%(message)s")
 logger = logging.getLogger("test-driver")
 logger.setLevel(logging.INFO)
 
+machine_colours_iter = (
+    "\x1b[{}m".format(x) for x in itertools.cycle(reversed(range(31, 37)))
+)
+
 
 class MachineLogAdapter(logging.LoggerAdapter):
     def process(self, msg: str, kwargs: Any) -> Tuple[str, Any]:
-        return f"{self.extra['machine']}: {msg}", kwargs
+        return (
+            f"{self.extra['colour_code']}{self.extra['machine']}\x1b[39m: {msg}",
+            kwargs,
+        )
 
 
 def make_command(args: list) -> str:
@@ -172,7 +180,10 @@ class Machine:
         self.socket = None
         self.monitor: Optional[socket.socket] = None
         self.allow_reboot = args.get("allowReboot", False)
-        self.logger = MachineLogAdapter(logger, extra=dict(machine=self.name))
+        self.logger = MachineLogAdapter(
+            logger,
+            extra=dict(machine=self.name, colour_code=next(machine_colours_iter)),
+        )
 
     @staticmethod
     def create_startcommand(args: Dict[str, str]) -> str: