share one hive-priv request helper across the vm tests

nix/vm-tests/agent-ephemeral-rootfs.nix

@@ -22,6 +22,7 @@
 # not booted here; the module-eval suite covers its wiring.
 { pkgs, hivePriv }:
 let
+  privRequest = import ./priv-request.nix { inherit pkgs; };
   agent = "iris";
   machine = "h-${agent}";
   agentDir = "/var/lib/hyperhive/agents/${agent}";
@@ -110,6 +111,8 @@ pkgs.testers.runNixOSTest {
   testScript = ''
     import json
 
+    ${privRequest}
+
     MACHINE = "${machine}"
     AGENT = "${agent}"
     AGENT_DIR = "${agentDir}"
@@ -135,24 +138,14 @@ pkgs.testers.runNixOSTest {
     def stop():
         host.succeed(f"nixos-container stop {MACHINE}")
 
-    def priv_journal():
-        return host.execute("journalctl -u hive-priv --no-pager | tail -n 30")[1]
-
     def write_flags():
         """What hive-c0re does before every start: one request to hive-priv."""
-        req = json.dumps({
+        priv_request(host, {
             "op": "write_nspawn_flags",
             "container": MACHINE,
             "binds": BINDS,
             "isolation": {"bridge": "${bridge}", "gateway_ip": "${gateway}"},
         })
-        # socat gives up 0.5s after its stdin ends unless told to wait; the
-        # salvage copies a whole root, which takes longer.
-        out = host.succeed(f"echo '{req}' | socat -t 120 - UNIX-CONNECT:/run/hive/priv.sock")
-        lines = out.strip().splitlines()
-        # The journal is read only when the assert fires.
-        assert lines, f"hive-priv sent no response; journal:\n{priv_journal()}"
-        assert json.loads(lines[-1])["ok"] is True, f"response={out!r}; journal:\n{priv_journal()}"
 
     def conf():
         return host.succeed(f"cat /etc/nixos-containers/{MACHINE}.conf")
@@ -241,12 +234,7 @@ pkgs.testers.runNixOSTest {
             {"op": "write_resource_limits", "container": MACHINE, "memory_max": "1G", "cpu_quota": "200%"},
             {"op": "daemon_reload"},
         ]:
-            # socat gives up 0.5s after its stdin ends unless told to wait, and
-            # a daemon-reload answers later than that.
-            out = host.succeed(f"echo '{json.dumps(req)}' | socat -t 30 - UNIX-CONNECT:/run/hive/priv.sock")
-            lines = out.strip().splitlines()
-            assert lines, f"hive-priv sent no response to {req['op']}; journal:\n{priv_journal()}"
-            assert json.loads(lines[-1])["ok"] is True, f"response={out!r}; journal:\n{priv_journal()}"
+            priv_request(host, req)
         dropin = host.succeed(f"cat /run/systemd/system/container@{MACHINE}.service.d/hyperhive-limits.conf")
         assert "ExecStartPre=/nix/store/" in dropin and f"check-rootfs {MACHINE}\n" in dropin, dropin
         # The build the drop-in names is rooted against garbage collection.

nix/vm-tests/priv-request.nix

@@ -0,0 +1,28 @@
+# Python for a VM test script: `priv_request`, the one way a test talks to
+# hive-priv. Splice it into `testScript` after `import json`. The VM it runs
+# against must have `socat` installed.
+{ pkgs }:
+''
+  import shlex
+
+  def priv_journal(node):
+      return node.execute("journalctl -u hive-priv --no-pager | tail -n 30")[1]
+
+  def priv_request(node, req, timeout=120):
+      """Send one request to hive-priv, wait for its reply, assert it is ok.
+
+      `-t` is how long socat waits for the reply after stdin ends; its default
+      is 0.5s, which drops the answer of any op that takes longer. Returns the
+      parsed reply.
+      """
+      op = req["op"]
+      out = node.succeed(
+          f"echo {shlex.quote(json.dumps(req))} | "
+          f"${pkgs.socat}/bin/socat -t {timeout} - UNIX-CONNECT:/run/hive/priv.sock"
+      )
+      lines = out.strip().splitlines()
+      assert lines, f"hive-priv sent no response to {op}; journal:\n{priv_journal(node)}"
+      reply = json.loads(lines[-1])
+      assert reply.get("ok") is True, f"{op}: response={out!r}; journal:\n{priv_journal(node)}"
+      return reply
+''