From d527c3334b35a886910948e41cea7db8a1b2a9bc Mon Sep 17 00:00:00 2001
From: Matthew Evans <git@ml-evs.science>
Date: Mon, 15 Jan 2024 22:46:11 +0000
Subject: [PATCH] Tweak pytest invocation

---
 .github/workflows/testing.yml   |  2 +-
 tests/integration/conftest.py   | 27 ++++++++++++++++++++++-----
 tests/integration/test_slurm.py | 10 ++++++----
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
index fe0622de..b9cf4894 100644
--- a/.github/workflows/testing.yml
+++ b/.github/workflows/testing.yml
@@ -58,7 +58,7 @@ jobs:
         run: pytest --cov=jobflow_remote --cov-report=xml --ignore tests/integration
 
       - name: Integration tests
-        run: pytest --cov=jobflow_remote --cov-append --cov-report=xml tests/integration
+        run: pytest --cov=jobflow_remote --cov-append --cov-report=xml tests/integration -s
 
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v3
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index 25c1cb7c..13da8be5 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -13,11 +13,28 @@
 from docker.models.containers import Container
 
 
-def _get_free_port():
-    """Returns a random free port."""
-    sock = socket.socket()
-    sock.bind(("", 0))
-    return sock.getsockname()[1]
+def _get_free_port(upper_bound=50_000):
+    """Returns a random free port, with an upper bound.
+
+    The upper bound is required as Docker does not have
+    permissions on high port numbers on some systems.
+
+    """
+    port = upper_bound + 1
+    attempts = 0
+    max_attempts = 10
+    while port > upper_bound and attempts < max_attempts:
+        sock = socket.socket()
+        sock.bind(("", 0))
+        port = sock.getsockname()[1]
+        attempts += 1
+
+    if attempts == max_attempts:
+        raise RuntimeError(
+            f"Could not find a free port to use with the provided {upper_bound=}."
+        )
+
+    return port
 
 
 def _get_random_name(length=6):
diff --git a/tests/integration/test_slurm.py b/tests/integration/test_slurm.py
index b114f16d..a87be8e1 100644
--- a/tests/integration/test_slurm.py
+++ b/tests/integration/test_slurm.py
@@ -36,6 +36,7 @@ def test_submit_flow(worker, job_controller):
 
     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
 
     add_first = add(1, 5)
@@ -47,9 +48,9 @@ def test_submit_flow(worker, job_controller):
     runner = Runner()
     runner.run(ticks=10)
 
-    assert job_controller.count_jobs({}) == 2
     assert len(job_controller.get_jobs({})) == 2
-    assert job_controller.count_flows({}) == 1
+    assert job_controller.count_jobs(state=JobState.COMPLETED) == 2
+    assert job_controller.count_flows(state=FlowState.COMPLETED) == 1
 
 
 @pytest.mark.parametrize(
@@ -61,6 +62,7 @@ def test_submit_flow_with_dependencies(worker, job_controller):
 
     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, write_file
 
     add_parent_1 = add(1, 1)
@@ -74,9 +76,9 @@ def test_submit_flow_with_dependencies(worker, job_controller):
     runner = Runner()
     runner.run(ticks=10)
 
-    assert job_controller.count_jobs({}) == 4
+    assert job_controller.count_jobs(state=JobState.COMPLETED) == 4
     assert len(job_controller.get_jobs({})) == 4
-    assert job_controller.count_flows({}) == 1
+    assert job_controller.count_flows(state=FlowState.COMPLETED) == 1
 
 
 @pytest.mark.parametrize(