From 1c5a031af36f27c7add692496bf1f35baa16f9bc Mon Sep 17 00:00:00 2001
From: Jared Baur <jaredbaur@fastmail.com>
Date: Mon, 22 Jul 2024 00:47:02 -0700
Subject: [PATCH] Make lrzsz work with ymodem test

---
 tests/ymodem/default.nix | 48 ++++++++++++++++++++++------------------
 1 file changed, 26 insertions(+), 22 deletions(-)

diff --git a/tests/ymodem/default.nix b/tests/ymodem/default.nix
index de4955f..2b1a812 100644
--- a/tests/ymodem/default.nix
+++ b/tests/ymodem/default.nix
@@ -31,27 +31,31 @@ testers.runNixOSTest {
       with open(params, "x") as f:
           f.write("init=${nodes.machine.system.build.toplevel}/init ${toString nodes.machine.boot.kernelParams}")
 
-      machine.start()
-      chardev = machine.send_monitor_command("chardev-add pty,id=ymodem")
-      machine.send_monitor_command("device_add virtconsole,chardev=ymodem")
-      machine.wait_for_console_text("press ENTER to interrupt")
-      machine.send_console("\n")  # interrupt boot process
-      time.sleep(1)
-      machine.send_console("list\n")
-      machine.send_console("select 2\n")  # TODO(jared): make this selection less vague
-      time.sleep(1)
-      machine.send_console("probe\n")
-
-      pty = re.findall(r"/dev/pts/[0-9]+", chardev)[0]
-      print(f"using pty {pty}")
-
-      subprocess.run(["${lib.getExe' tinybootTools "tboot-ymodem"}", "send", "--tty", pty, "--dir", host_boot_dir.name])
-      # with open(pty, "r+b") as f:
-      #     TODO(jared): make this work with stdin/stdout
-      #     subprocess.run(["${lib.getExe' lrzsz "sx"}", "--ymodem", "-kb", linux, initrd, params], stdin=f, stdout=f, text=False)
-
-      machine.send_console("boot\n")
-
-      machine.wait_for_unit("multi-user.target")
+      def tboot_ymodem(pty):
+          subprocess.run(["${lib.getExe' tinybootTools "tboot-ymodem"}", "send", "--tty", pty, "--dir", host_boot_dir.name])
+
+      def lrzsz(pty):
+          subprocess.run(f"${lib.getExe' lrzsz "sx"} --ymodem -kb {linux} {initrd} {params} > {pty} < {pty}", shell=True)
+
+      for fn in [tboot_ymodem, lrzsz]:
+          machine.start()
+          chardev = machine.send_monitor_command("chardev-add pty,id=ymodem")
+          machine.send_monitor_command("device_add virtconsole,chardev=ymodem")
+          machine.wait_for_console_text("press ENTER to interrupt")
+          machine.send_console("\n")  # interrupt boot process
+          time.sleep(1)
+          machine.send_console("list\n")
+          machine.send_console("select 2\n")  # TODO(jared): make this selection less vague
+          time.sleep(1)
+          machine.send_console("probe\n")
+
+          pty = re.findall(r"/dev/pts/[0-9]+", chardev)[0]
+          print(f"using pty {pty}")
+          fn(pty)
+
+          machine.send_console("boot\n")
+          machine.wait_for_unit("multi-user.target")
+          machine.shutdown()
+          machine.wait_for_shutdown()
     '';
 }