From fa38416f85b0e3a9443bd873d1b7d6d00d5d196f Mon Sep 17 00:00:00 2001 From: Matthew Evans Date: Fri, 26 Jan 2024 13:33:56 +0000 Subject: [PATCH] Tweak arithmetic/callable test --- src/jobflow_remote/testing/__init__.py | 13 +++++++++---- tests/integration/test_slurm.py | 2 +- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/jobflow_remote/testing/__init__.py b/src/jobflow_remote/testing/__init__.py index 961e1166..581dd9da 100644 --- a/src/jobflow_remote/testing/__init__.py +++ b/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 diff --git a/tests/integration/test_slurm.py b/tests/integration/test_slurm.py index 066e0417..d7ae29ad 100644 --- a/tests/integration/test_slurm.py +++ b/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])