From d50a8aae11ad395e282e28ea89e1be4f89470b61 Mon Sep 17 00:00:00 2001
From: Matthew Evans <git@ml-evs.science>
Date: Mon, 29 Jan 2024 12:25:02 +0000
Subject: [PATCH] Add a test that runs with an exec_config and checks an
 environment variable is set in the job

---
 src/jobflow_remote/testing/__init__.py |  7 +++++++
 tests/integration/conftest.py          |  1 +
 tests/integration/test_slurm.py        | 29 ++++++++++++++++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/src/jobflow_remote/testing/__init__.py b/src/jobflow_remote/testing/__init__.py
index cf3eec46..4aafe91e 100644
--- a/src/jobflow_remote/testing/__init__.py
+++ b/src/jobflow_remote/testing/__init__.py
@@ -34,3 +34,10 @@ def arithmetic(
         return op(a, b)
 
     return None
+
+
+@job
+def check_env_var() -> str:
+    import os
+
+    return os.environ.get("TESTING_ENV_VAR", "unset")
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index 5163cd95..903d0415 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -229,6 +229,7 @@ def write_tmp_settings(
                 connect_kwargs={"allow_agent": False, "look_for_keys": False},
             ),
         },
+        exec_config={"test": {"export": {"TESTING_ENV_VAR": random_project_name}}},
         runner=dict(
             delay_checkout=1,
             delay_check_run_status=1,
diff --git a/tests/integration/test_slurm.py b/tests/integration/test_slurm.py
index 9c935ff3..7dc849e9 100644
--- a/tests/integration/test_slurm.py
+++ b/tests/integration/test_slurm.py
@@ -205,3 +205,32 @@ def test_expected_failure(worker, job_controller):
 
     assert job_controller.count_jobs(state=JobState.FAILED) == 2
     assert job_controller.count_flows(state=FlowState.FAILED) == 1
+
+
+@pytest.mark.parametrize(
+    "worker",
+    ["test_local_worker", "test_remote_worker"],
+)
+def test_exec_config(worker, job_controller, random_project_name):
+    """Tests that an environment variable set in the exec config
+    is available to the job.
+
+    """
+
+    from jobflow_remote import submit_flow
+    from jobflow_remote.jobs.runner import Runner
+    from jobflow_remote.testing import check_env_var
+
+    job = check_env_var()
+    submit_flow(job, worker=worker, exec_config="test")
+
+    assert job_controller.count_jobs({}) == 1
+    assert len(job_controller.get_jobs({})) == 1
+    assert job_controller.count_flows({}) == 1
+
+    runner = Runner()
+    runner.run(ticks=5)
+
+    job = job_controller.get_jobs({})[0]
+    output = job_controller.jobstore.get_output(uuid=job["uuid"])
+    assert output == random_project_name