nixosTest: remove hostname limitations

2 files changed, 17 insertions, 21 deletions

diff --git a/nixos/lib/test-driver/test_driver/driver.py b/nixos/lib/test-driver/test_driver/driver.py
index ad52f365737c..ea6ba4b65b56 100644
--- a/nixos/lib/test-driver/test_driver/driver.py
+++ b/nixos/lib/test-driver/test_driver/driver.py
@@ -2,6 +2,7 @@ from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Union, Optional, Callable, ContextManager
 import os
+import re
 import tempfile
 
 from test_driver.logger import rootlog
@@ -28,6 +29,10 @@ def get_tmp_dir() -> Path:
     return tmp_dir
 
 
+def pythonize_name(name: str) -> str:
+    return re.sub(r"^[^A-z_]|[^A-z0-9_]", "_", name)
+
+
 class Driver:
     """A handle to the driver that sets up the environment
     and runs the tests"""
@@ -113,7 +118,7 @@ class Driver:
             polling_condition=self.polling_condition,
             Machine=Machine,  # for typing
         )
-        machine_symbols = {m.name: m for m in self.machines}
+        machine_symbols = {pythonize_name(m.name): m for m in self.machines}
         # If there's exactly one machine, make it available under the name
         # "machine", even if it's not called that.
         if len(self.machines) == 1:
diff --git a/nixos/lib/testing/driver.nix b/nixos/lib/testing/driver.nix
index fb181c1d7e9a..25759a91dda3 100644
--- a/nixos/lib/testing/driver.nix
+++ b/nixos/lib/testing/driver.nix
@@ -21,29 +21,20 @@ let
     in
     nodesList ++ lib.optional (lib.length nodesList == 1 && !lib.elem "machine" nodesList) "machine";
 
-  # TODO: This is an implementation error and needs fixing
-  # the testing famework cannot legitimately restrict hostnames further
-  # beyond RFC1035
-  invalidNodeNames = lib.filter
-    (node: builtins.match "^[A-z_]([A-z0-9_]+)?$" node == null)
-    nodeHostNames;
+  pythonizeName = name:
+    let
+      head = lib.substring 0 1 name;
+      tail = lib.substring 1 (-1) name;
+    in
+      (if builtins.match "[A-z_]" head == null then "_" else head) +
+      lib.stringAsChars (c: if builtins.match "[A-z0-9_]" c == null then "_" else c) tail;
 
   uniqueVlans = lib.unique (builtins.concatLists vlans);
   vlanNames = map (i: "vlan${toString i}: VLan;") uniqueVlans;
-  machineNames = map (name: "${name}: Machine;") nodeHostNames;
+  pythonizedNames = map pythonizeName nodeHostNames;
+  machineNames = map (name: "${name}: Machine;") pythonizedNames;
 
-  withChecks =
-    if lib.length invalidNodeNames > 0 then
-      throw ''
-        Cannot create machines out of (${lib.concatStringsSep ", " invalidNodeNames})!
-        All machines are referenced as python variables in the testing framework which will break the
-        script when special characters are used.
-
-        This is an IMPLEMENTATION ERROR and needs to be fixed. Meanwhile,
-        please stick to alphanumeric chars and underscores as separation.
-      ''
-    else
-      lib.warnIf config.skipLint "Linting is disabled";
+  withChecks = lib.warnIf config.skipLint "Linting is disabled";
 
   driver =
     hostPkgs.runCommand "nixos-test-driver-${config.name}"
@@ -87,7 +78,7 @@ let
         ${testDriver}/bin/generate-driver-symbols
         ${lib.optionalString (!config.skipLint) ''
           PYFLAKES_BUILTINS="$(
-            echo -n ${lib.escapeShellArg (lib.concatStringsSep "," nodeHostNames)},
+            echo -n ${lib.escapeShellArg (lib.concatStringsSep "," pythonizedNames)},
             < ${lib.escapeShellArg "driver-symbols"}
           )" ${hostPkgs.python3Packages.pyflakes}/bin/pyflakes $out/test-script
         ''}