Added an optional on_exit to the autotest stressor library.
BUG=None
TEST=stress_unittest.py
Change-Id: Ia12be63c96114076b85ecdad0e6def064969fd60
Reviewed-on: https://gerrit.chromium.org/gerrit/38430
Reviewed-by: Prashanth Balasubramanian <beeps@chromium.org>
Reviewed-by: Kris Rambish <krisr@chromium.org>
Commit-Ready: Craig Harrison <craigdh@chromium.org>
Tested-by: Craig Harrison <craigdh@chromium.org>
diff --git a/server/cros/stress.py b/server/cros/stress.py
index 2b0d525..10495c9 100755
--- a/server/cros/stress.py
+++ b/server/cros/stress.py
@@ -14,17 +14,19 @@
@var stressor: callable which performs a single stress event.
"""
- def __init__(self, stressor, escalate_exceptions=True):
+ def __init__(self, stressor, on_exit=None, escalate_exceptions=True):
"""
Initialize the ControlledStressor.
@param stressor: callable which performs a single stress event.
+ @param on_exit: callable which will be called when the thread finishes.
@param escalate_exceptions: whether to escalate exceptions to the parent
thread; defaults to True.
"""
super(BaseStressor, self).__init__()
self.daemon = True
self.stressor = stressor
+ self.on_exit = on_exit
self._escalate_exceptions = escalate_exceptions
self._exc_info = None
@@ -59,6 +61,9 @@
if self._escalate_exceptions:
self._exc_info = sys.exc_info()
raise
+ finally:
+ if self.on_exit:
+ self.on_exit()
def _loop_stressor(self):
@@ -88,16 +93,18 @@
Creates a new thread and calls |stressor| in a loop until stop() is called.
"""
- def __init__(self, stressor, escalate_exceptions=True):
+ def __init__(self, stressor, on_exit=None, escalate_exceptions=True):
"""
Initialize the ControlledStressor.
@param stressor: callable which performs a single stress event.
+ @param on_exit: callable which will be called when the thread finishes.
@param escalate_exceptions: whether to escalate exceptions to the parent
thread; defaults to True.
"""
self._complete = threading.Event()
- super(ControlledStressor, self).__init__(stressor, escalate_exceptions)
+ super(ControlledStressor, self).__init__(stressor, on_exit,
+ escalate_exceptions)
def _loop_stressor(self):
diff --git a/server/cros/stress_unittest.py b/server/cros/stress_unittest.py
index 6ddcd90..44e3091 100644
--- a/server/cros/stress_unittest.py
+++ b/server/cros/stress_unittest.py
@@ -36,6 +36,34 @@
self.assertTrue(event.is_set(), 'The stress event did not run')
+ def testOnExit(self):
+ def stress_event():
+ pass
+
+ event = threading.Event()
+ def on_exit():
+ event.set()
+
+ stressor = stress.CountedStressor(stress_event, on_exit=on_exit)
+ stressor.start(1)
+ stressor.wait()
+ self.assertTrue(event.is_set())
+
+
+ def testOnExitWithException(self):
+ def stress_event():
+ raise StopThreadForTesting
+
+ event = threading.Event()
+ def on_exit():
+ event.set()
+
+ stressor = stress.CountedStressor(stress_event, on_exit=on_exit)
+ stressor.start(1)
+ self.assertRaises(StopThreadForTesting, stressor.wait)
+ self.assertTrue(event.is_set())
+
+
def testCountedStressorStartCondition(self):
event = threading.Event()