Fixed Python tests

spirali · spirali · commit 4c9a5e4f382c · 2025-05-06T19:20:12.000+02:00
diff --git a/tests/test_journal.py b/tests/test_journal.py
@@ -348,7 +348,7 @@ def collect_ids():
 def test_restore_dependencies1(hq_env: HqEnv, tmp_path):
     journal_path = os.path.join(tmp_path, "my.journal")
     hq_env.start_server(args=["--journal", journal_path])
-    hq_env.start_worker()
+    hq_env.start_worker(cpus=2)
     tmp_path.joinpath("sleep_time").write_text("100")
     tmp_path.joinpath("job.toml").write_text(
         """
diff --git a/tests/test_worker.py b/tests/test_worker.py
@@ -317,7 +317,7 @@ def test_deploy_ssh_error(hq_env: HqEnv):
     with hq_env.mock.mock_program_with_code(
         "ssh",
         """
-                exit(42)
+        exit(42)
         """,
     ):
         nodefile = prepare_localhost_nodefile()
@@ -351,8 +351,8 @@ def test_deploy_ssh_wait_for_worker(hq_env: HqEnv):
     with hq_env.mock.mock_program_with_code(
         "ssh",
         """
-                    import time
-                    time.sleep(1)
+        import time
+        time.sleep(1)
         """,
     ):
         nodefile = prepare_localhost_nodefile()
@@ -364,12 +364,12 @@ def test_deploy_ssh_multiple_workers(hq_env: HqEnv):
     with hq_env.mock.mock_program_with_code(
         "ssh",
         """
-                    import os
+        import os
 
-                    os.makedirs("workers", exist_ok=True)
-                    with open(f"workers/{os.getpid()}", "w") as f:
-                        f.write(str(os.getpid()))
-                    """,
+        os.makedirs("workers", exist_ok=True)
+        with open(f"workers/{os.getpid()}", "w") as f:
+            f.write(str(os.getpid()))
+        """,
     ):
         nodefile = prepare_localhost_nodefile(count=3)
         hq_env.command(["worker", "deploy-ssh", nodefile])
@@ -381,8 +381,8 @@ def test_deploy_ssh_show_output(hq_env: HqEnv):
     with hq_env.mock.mock_program_with_code(
         "ssh",
         """
-                    print("FOOBAR")
-            """,
+        print("FOOBAR")
+        """,
     ):
         nodefile = prepare_localhost_nodefile(count=3)
         output = hq_env.command(["worker", "deploy-ssh", "--show-output", nodefile])
@@ -424,15 +424,15 @@ def test_worker_state_info(hq_env: HqEnv):
     wait_for_job_state(hq_env, 1, "RUNNING")
     table = hq_env.command(["worker", "info", "1"], as_table=True)
     table.check_row_value("State", "RUNNING")
-    assert "Job: 1;" in table.get_row_value("Last task started")
+    assert "1@" in table.get_row_value("Last task started")
     table.check_row_value("Runtime Info", "assigned tasks: 2; running tasks: 1")
     hq_env.start_worker()
     hq_env.command(["submit", "--nodes=2", "--", "sleep", "1"])
     wait_for_job_state(hq_env, 2, "RUNNING")
     table = hq_env.command(["worker", "info", "1"], as_table=True)
     a = table.get_row_value("Runtime Info")
-    assert "Job: 2;" in table.get_row_value("Last task started")
+    assert "2@0" in table.get_row_value("Last task started")
     table = hq_env.command(["worker", "info", "2"], as_table=True)
     b = table.get_row_value("Runtime Info")
-    assert "Job: 2;" in table.get_row_value("Last task started")
+    assert "2@0" in table.get_row_value("Last task started")
     assert {a, b} == {"running multinode task; main node", "running multinode task; secondary node"}
diff --git a/tests/utils/wait.py b/tests/utils/wait.py
@@ -90,12 +90,17 @@ def wait_for_task_state(env, job_id: int, task_ids: Union[int, List[int]], state
 
     ids = ",".join(str(t) for t in task_ids)
     states = [s.lower() for s in states]
+    result = None
 
     def check():
+        nonlocal result
         result = env.command(["--output-mode=json", "task", "info", str(job_id), ids], as_json=True)
         return [r["state"] for r in result] == states
 
-    wait_until(check, **kwargs)
+    def on_timeout():
+        return f"most recent output:\n{result}"
+
+    wait_until(check, on_timeout=on_timeout, **kwargs)
 
 
 def wait_for_pid_exit(pid: int):
