added tests and made fixes

Author: Guy Amram

computation_graph/graph_runners.py

@@ -1,4 +1,4 @@
-import inspect
+import asyncio
 from typing import Callable, Union
 
 import gamla
@@ -78,16 +78,22 @@ def unary_with_state_and_expectations(
 def variadic_with_state_and_expectations(g, sink):
     f = run.to_callable_strict(g)
 
-    async def inner(turns):
-        prev = {}
-        for turn, expectation in turns:
-            if inspect.iscoroutinefunction(f):
+    if asyncio.iscoroutinefunction(f):
+        async def inner(turns):
+            prev = {}
+            for turn, expectation in turns:
                 prev = await f(prev, turn)
-            else:
+                assert (
+                    prev[graph.make_computation_node(sink)] == expectation
+                ), f"actual={prev[graph.make_computation_node(sink)]}\n expected: {expectation}"
+    else:
+        def inner(turns):
+            prev = {}
+            for turn, expectation in turns:
                 prev = f(prev, turn)
-            assert (
-                prev[graph.make_computation_node(sink)] == expectation
-            ), f"actual={prev[graph.make_computation_node(sink)]}\n expected: {expectation}"
+                assert (
+                    prev[graph.make_computation_node(sink)] == expectation
+                ), f"actual={prev[graph.make_computation_node(sink)]}\n expected: {expectation}"
 
     return inner
 

computation_graph/graph_test.py

@@ -499,6 +499,11 @@ def _plus_1(y):
     return y + 1
 
 
+async def _plus_1_async(y):
+    await asyncio.sleep(0.1)
+    return y + 1
+
+
 def _times_2(x):
     return x * 2
 
@@ -576,6 +581,26 @@ def test_compose_future():
     )(([[{a: 2, b: 2, c: 2}, 9], [{a: 2, b: 2, c: 2}, 25]]))
 
 
+def test_compose_future_async():
+    a = graph.make_source()
+    b = graph.make_source()
+    c = graph.make_source()
+    graph_runners.variadic_with_state_and_expectations(
+        base_types.merge_graphs(
+            composers.compose_source_unary(_plus_1_async, c),
+            composers.compose_source_unary(_times_2, b),
+            composers.compose_source(_multiply, "a", a),
+            composers.make_compose_future(
+                _multiply,
+                composers.make_and([_plus_1_async, _times_2, _multiply], merge_fn=_sum),
+                "b",
+                None,
+            ),
+        ),
+        _sum,
+    )(([[{a: 2, b: 2, c: 2}, 9], [{a: 2, b: 2, c: 2}, 25]]))
+
+
 def test_dont_duplicate_sources():
     a = graph.make_source()
     assert (