Make scenario finalizing more functional

README.md

@@ -109,20 +109,20 @@ You can also launch multiple multi-UA scenarios concurrently using
 non-blocking mode:
 
 ```python
-scens = []
+finalizers = []
 for path, scen in pysipp.walk('path/to/scendirs/root/', proxyaddr=('10.10.8.1', 5060)):
     print("Running scenario collected from {}".format(path))
-    scen(block=False)
-    scens.append(scen)
+    finalizers.append(scen(block=False))
 
 # All scenarios should now be running concurrently so we can continue
 # program execution...
 print("Continuing program execution...")
 
-# Wait to collect all the results
+# Calling `scen` returns a finalizer that can be
+# used to wait for each scenario with a timeout
 print("Finalizing all scenarios and collecting results")
-for scen in scens:
-  scen.finalize()
+for finalizer in finalizers:
+  finalizer(timeout=10)
 ```
 
 `scen.finalize()` actually calls a special cleanup function defined in the

pysipp/agent.py

@@ -441,25 +441,14 @@ async def arun(
         runner=None,
         block=True,
     ):
-        self._prepared_agents = agents = self.prepare()
-        self._runner = runner = runner or launch.TrioRunner()
-
-        return await launch.run_all_agents(runner, agents, timeout=timeout, block=block)
-
-    def finalize(self, *, timeout=DEFAULT_RUNNER_TIMEOUT):
-        assert (
-            self._prepared_agents and self._runner
-        ), "Must run scenario before finalizing."
-        return trio.run(
-            partial(
-                launch.finalize,
-                self._runner,
-                self._prepared_agents,
-                timeout=timeout,
-            )
+        agents = self.prepare()
+        self.runner = runner or launch.TrioRunner()
+
+        return await launch.run_all_agents(
+            self.runner, agents, timeout=timeout, block=block
         )
 
-    def run(self, timeout=180, **kwargs):
+    def run(self, timeout=DEFAULT_RUNNER_TIMEOUT, **kwargs):
         """Run scenario blocking to completion."""
         return trio.run(partial(self.arun, timeout=timeout, **kwargs))
 

pysipp/launch.py

@@ -6,8 +6,9 @@
 import subprocess
 import time
 from . import utils
-from pprint import pformat
 from collections import OrderedDict, namedtuple
+from functools import partial
+from pprint import pformat
 
 import trio
 
@@ -149,18 +150,19 @@ def clear(self):
         self._procs.clear()
 
 
-async def run_all_agents(runner, agents, timeout=180, block=True):
+async def run_all_agents(runner, agents, timeout, block=True):
     """Run a sequencec of agents using a ``TrioRunner``."""
 
     try:
         await runner.run((ua.render() for ua in agents), timeout=timeout)
         if block:
-            await finalize(runner, agents, timeout)
-        return runner
+            return await finalize(runner, agents, timeout)
+        else:
+            return finalizer(runner, agents)
     except TimeoutError as terr:
         # print error logs even when we timeout
         try:
-            await finalize(runner, agents, timeout)
+            return await finalize(runner, agents, timeout)
         except SIPpFailure as err:
             assert "exit code -9" in str(err)
             raise terr
@@ -178,3 +180,10 @@ async def finalize(runner, agents, timeout):
         raise SIPpFailure(msg)
 
     return cmds2procs
+
+
+def finalizer(runner, agents):
+    def with_timeout(timeout):
+        return trio.run(partial(finalize, runner, agents, timeout=timeout))
+
+    return with_timeout

tests/test_launcher.py

@@ -10,7 +10,7 @@
 
 def run_blocking(runner, agents):
     assert not runner.is_alive()
-    trio.run(run_all_agents, runner, agents)
+    trio.run(run_all_agents, runner, agents, 10)
     assert not runner.is_alive()
     return runner
 

tests/test_stack.py

@@ -73,29 +73,29 @@ def test_proxyaddr_with_scendir(scendir):
 def test_sync_run(scenwalk):
     """Ensure all scenarios in the test run to completion in synchronous mode"""
     for path, scen in scenwalk():
-        runner = scen(timeout=6)
-        for cmd, proc in runner._procs.items():
+        finalizer = scen(timeout=6)
+        for cmd, proc in finalizer.items():
             assert proc.returncode == 0
 
 
 def test_async_run(scenwalk):
     """Ensure multiple scenarios run to completion in asynchronous mode."""
     finalizers = []
-    for path, scen in scenwalk():
-        finalizers.append((scen, scen(block=False)))
+    for _, scen in scenwalk():
+        finalizers.append(scen(block=False))
 
     # collect all results synchronously
-    for scen, finalizer in finalizers:
-        for cmd, proc in scen.finalize(timeout=6).items():
+    for finalizer in finalizers:
+        for cmd, proc in finalizer(timeout=6).items():
             assert proc.returncode == 0
 
 
 def test_basic(basic_scen):
     """Test the most basic uac <-> uas call flow"""
     assert len(basic_scen.agents) == 2
     # ensure sync run works
-    runner = basic_scen()
-    assert not runner.is_alive()
+    basic_scen()
+    assert not basic_scen.runner.is_alive()
 
 
 def test_unreachable_uas(basic_scen):