diff --git a/catkit/testbed/experiment.py b/catkit/testbed/experiment.py index 1f2c1fb0..44ce5008 100644 --- a/catkit/testbed/experiment.py +++ b/catkit/testbed/experiment.py @@ -94,7 +94,7 @@ def __init__(self, safety_tests, output_path=None, suffix=None, for test in safety_tests: self.safety_tests.append(test()) - def __enter__(self): + def start(self): try: self._setup() @@ -125,9 +125,15 @@ def __enter__(self): finally: raise + def stop(self): + return self._teardown() + + def __enter__(self): + return self.start() + def __exit__(self, exc_type, exc_val, exc_tb): try: - self._teardown() + self.stop() finally: if exc_type: self.log.exception("The testbed encountered the following error(s):")