From 4abf99b2031a3f315688b870c853b7bffab70ca9 Mon Sep 17 00:00:00 2001
From: Matthew Evans <git@ml-evs.science>
Date: Sun, 18 Feb 2024 15:01:49 +0000
Subject: [PATCH] Add test for case where additional store is not known to
 jobflow-remote

---
 src/jobflow_remote/testing/__init__.py | 15 ++++++++++++-
 tests/integration/test_slurm.py        | 29 ++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/jobflow_remote/testing/__init__.py b/src/jobflow_remote/testing/__init__.py
index 53d40963..e3d6362a 100644
--- a/src/jobflow_remote/testing/__init__.py
+++ b/src/jobflow_remote/testing/__init__.py
@@ -45,7 +45,20 @@ def check_env_var() -> str:
 
 @job(big_files="data")
 def add_big(a: float, b: float):
-    """Adds two numbers together and writes the answer to an artificially large file."""
+    """Adds two numbers together and writes the answer to an artificially large file
+    which is stored in a pre-defined store."""
+    import pathlib
+
+    result = a + b
+    with open("file.txt", "w") as f:
+        f.writelines([f"{result}"] * int(1e5))
+    return Response({"data": pathlib.Path("file.txt"), "result": a + b})
+
+
+@job(undefined_store="data")
+def add_big_undefined_store(a: float, b: float):
+    """Adds two numbers together and writes the answer to an artificially large file
+    which is attempted to be stored in a undefined store."""
     import pathlib
 
     result = a + b
diff --git a/tests/integration/test_slurm.py b/tests/integration/test_slurm.py
index 41eed645..10d13c14 100644
--- a/tests/integration/test_slurm.py
+++ b/tests/integration/test_slurm.py
@@ -268,3 +268,32 @@ def test_additional_stores(worker, job_controller):
         "blob_uuid"
     ]
     assert fs.query({"blob_uuid": blob_uuid})[0]["job_uuid"] == doc["job"]["uuid"]
+
+
+@pytest.mark.parametrize(
+    "worker",
+    ["test_local_worker", "test_remote_worker"],
+)
+def test_undefined_additional_stores(worker, job_controller):
+    from jobflow import Flow
+
+    from jobflow_remote import submit_flow
+    from jobflow_remote.jobs.runner import Runner
+    from jobflow_remote.jobs.state import FlowState, JobState
+    from jobflow_remote.testing import add_big_undefined_store
+
+    job = add_big_undefined_store(100, 100)
+    flow = Flow(job)
+    submit_flow(flow, worker=worker)
+
+    assert job_controller.count_jobs({}) == 1
+    assert job_controller.count_flows({}) == 1
+
+    runner = Runner()
+    runner.run(ticks=10)
+
+    # Probably this should error somewhere
+    doc = job_controller.get_jobs({})[0]
+    assert job_controller.count_jobs(state=JobState.COMPLETED) == 1
+    assert job_controller.count_flows(state=FlowState.COMPLETED) == 1
+    assert job_controller.jobstore.get_output(uuid=doc["job"]["uuid"])["result"] == 200