Tweak arithmetic/callable test

src/jobflow_remote/testing/__init__.py

@@ -1,5 +1,5 @@
 """A series of toy workflows that can be used for testing."""
-from typing import Callable
+from typing import Callable, Optional
 
 from jobflow import job
 
@@ -25,6 +25,11 @@ def write_file(n):
 
 @job
 def arithmetic(
-    a: float, b: float, op: Callable[[float, float], float] = lambda x, y: x + y
-) -> float:
-    return op(a, b)
+    a,
+    b,
+    op: Optional[Callable] = None,
+) -> Optional[float]:
+    if op:
+        return op(a, b)
+
+    return None

tests/integration/test_slurm.py

@@ -153,7 +153,7 @@ def test_job_with_callable_kwarg(worker, job_controller):
     from jobflow_remote.testing import arithmetic
 
     job_1 = arithmetic(1, -2, op=math.copysign)
-    job_2 = arithmetic(job_1.output, 1, op=math.dist)
+    job_2 = arithmetic([job_1.output], [1], op=math.dist)
     job_3 = arithmetic(job_2.output, 2, op=math.pow)
 
     flow = Flow([job_1, job_2, job_3])