init

tests/test_session.py

@@ -0,0 +1,54 @@
+import os
+
+from smallpond.dataframe import Session
+
+
+def test_shutdown_cleanup(sp: Session):
+    assert os.path.exists(sp._runtime_ctx.queue_root), "queue directory should exist"
+    assert os.path.exists(
+        sp._runtime_ctx.staging_root
+    ), "staging directory should exist"
+    assert os.path.exists(sp._runtime_ctx.temp_root), "temp directory should exist"
+
+    # create some tasks and complete them
+    df = sp.from_items([1, 2, 3])
+    df.write_parquet(sp._runtime_ctx.output_root)
+    sp.shutdown()
+
+    # shutdown should clean up directories
+    assert not os.path.exists(
+        sp._runtime_ctx.queue_root
+    ), "queue directory should be cleared"
+    assert not os.path.exists(
+        sp._runtime_ctx.staging_root
+    ), "staging directory should be cleared"
+    assert not os.path.exists(
+        sp._runtime_ctx.temp_root
+    ), "temp directory should be cleared"
+    with open(sp._runtime_ctx.job_status_path) as fin:
+        assert "success" in fin.read(), "job status should be success"
+
+
+def test_shutdown_no_cleanup_on_failure(sp: Session):
+    df = sp.from_items([1, 2, 3])
+    try:
+        # create a task that will fail
+        df.map(lambda x: x / 0).compute()
+    except Exception:
+        pass
+    else:
+        raise RuntimeError("task should fail")
+    sp.shutdown()
+
+    # shutdown should not clean up directories
+    assert os.path.exists(
+        sp._runtime_ctx.queue_root
+    ), "queue directory should not be cleared"
+    assert os.path.exists(
+        sp._runtime_ctx.staging_root
+    ), "staging directory should not be cleared"
+    assert os.path.exists(
+        sp._runtime_ctx.temp_root
+    ), "temp directory should not be cleared"
+    with open(sp._runtime_ctx.job_status_path) as fin:
+        assert "failure" in fin.read(), "job status should be failure"