Test for server restoration & streaming

spirali · spirali · commit 8796f11595f4 · 2024-09-20T13:31:37.000+02:00
diff --git a/tests/test_restore.py b/tests/test_restore.py
@@ -1,5 +1,6 @@
 import time
 
+from .utils.wait import wait_until
 from .utils.cmd import python
 from .autoalloc.mock.mock import MockJobManager
 from .autoalloc.mock.slurm import SlurmManager, adapt_slurm
@@ -270,3 +271,28 @@ def test_restore_open_job_more_submits(hq_env: HqEnv, tmp_path):
     table.check_row_value("Tasks", "3; Ids: 0-2")
     hq_env.start_worker()
     wait_for_job_state(hq_env, 1, "OPENED")
+
+
+def test_restore_streaming(hq_env: HqEnv, tmp_path):
+    journal_path = os.path.join(tmp_path, "my.journal")
+    stream_path = os.path.join(tmp_path, "stream")
+    os.mkdir(stream_path)
+    hq_env.start_server(args=["--journal", journal_path])
+    hq_env.start_worker()
+    hq_env.command(["job", "open"])
+    hq_env.command(
+        ["job", "submit", "--job=1", "--stream", stream_path, "--", "bash", "-c", "echo $HQ_INSTANCE_ID; sleep 2"]
+    )
+    wait_for_job_state(hq_env, 1, "RUNNING")
+    hq_env.stop_server()
+    hq_env.start_server(args=["--journal", journal_path])
+
+    table = hq_env.command(["job", "info", "1"], as_table=True)
+    table.check_row_value("Stdout", "<Stream>")
+    table.check_row_value("Stderr", "<Stream>")
+    hq_env.start_worker()
+    hq_env.command(["job", "close", "1"])
+    wait_for_job_state(hq_env, 1, "FINISHED")
+    assert int(hq_env.command(["read", stream_path, "cat", "1", "stdout"])) > 0
+    table = hq_env.command(["read", stream_path, "summary"], as_table=True)
+    table.check_row_value("Superseded streams", "1")
diff --git a/tests/test_stream.py b/tests/test_stream.py
@@ -100,7 +100,7 @@ def test_stream_overlap(hq_env: HqEnv):
     wait_for_job_state(hq_env, 1, "FINISHED")
     result = hq_env.command(["read", "mylog", "show"], as_lines=True)
 
-    chunks = [set(result[i * 2: i * 2 + 2]) for i in range(3)]
+    chunks = [set(result[i * 2 : i * 2 + 2]) for i in range(3)]
     assert chunks[0] == {"1.1:0> A", "1.2:0> A"}
     assert chunks[1] == {"1.2:0> B", "1.1:0> B"}
 
