Remove mox from test_scheduler.TaskTest

Change-Id: Iff1b1632c66d6727798969a8ca6506d8d0fb2051

1 files changed, 244 insertions, 332 deletions

diff --git a/heat/tests/engine/test_scheduler.py b/heat/tests/engine/test_scheduler.py
index 386e51891..017a6989b 100644
--- a/heat/tests/engine/test_scheduler.py
+++ b/heat/tests/engine/test_scheduler.py
@@ -428,34 +428,23 @@ class TaskTest(common.HeatTestCase):
     def setUp(self):
         super(TaskTest, self).setUp()
         scheduler.ENABLE_SLEEP = True
-        self.addCleanup(self.m.VerifyAll)
+        self.mock_sleep = self.patchobject(scheduler.TaskRunner, '_sleep',
+                                           return_value=None)
 
     def test_run(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(0).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         scheduler.TaskRunner(task)()
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.assertEqual(3, self.mock_sleep.call_count)
+
     def test_run_as_task(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         tr = scheduler.TaskRunner(task)
         rt = tr.as_task()
@@ -463,16 +452,14 @@ class TaskTest(common.HeatTestCase):
             pass
         self.assertTrue(tr.done())
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+
     def test_run_as_task_started(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         tr = scheduler.TaskRunner(task)
         tr.start()
@@ -480,14 +467,14 @@ class TaskTest(common.HeatTestCase):
             pass
         self.assertTrue(tr.done())
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+
     def test_run_as_task_cancel(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         tr = scheduler.TaskRunner(task)
         rt = tr.as_task()
@@ -496,17 +483,15 @@ class TaskTest(common.HeatTestCase):
 
         self.assertTrue(tr.done())
 
+        task.do_step.assert_called_once_with(1)
+        self.mock_sleep.assert_not_called()
+
     def test_run_as_task_exception(self):
         class TestException(Exception):
             pass
 
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         tr = scheduler.TaskRunner(task)
         rt = tr.as_task()
@@ -515,6 +500,9 @@ class TaskTest(common.HeatTestCase):
 
         self.assertTrue(tr.done())
 
+        task.do_step.assert_called_once_with(1)
+        self.mock_sleep.assert_not_called()
+
     def test_run_as_task_swallow_exception(self):
         class TestException(Exception):
             pass
@@ -536,93 +524,81 @@ class TaskTest(common.HeatTestCase):
 
     def test_run_delays(self):
         task = DummyTask(delays=itertools.repeat(2))
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(0).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         scheduler.TaskRunner(task)()
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(0),
+                                          mock.call(1),
+                                          mock.call(1),
+                                          mock.call(1),
+                                          mock.call(1),
+                                          mock.call(1)])
+        self.assertEqual(6, self.mock_sleep.call_count)
+
     def test_run_delays_dynamic(self):
         task = DummyTask(delays=[2, 4, 1])
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(0).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         scheduler.TaskRunner(task)()
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(0),
+                                          mock.call(1),
+                                          mock.call(1),
+                                          mock.call(1),
+                                          mock.call(1),
+                                          mock.call(1),
+                                          mock.call(1)])
+        self.assertEqual(7, self.mock_sleep.call_count)
+
     def test_run_wait_time(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(0).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(42).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-        scheduler.TaskRunner._sleep(42).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         scheduler.TaskRunner(task)(wait_time=42)
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(0),
+                                          mock.call(42),
+                                          mock.call(42)])
+        self.assertEqual(3, self.mock_sleep.call_count)
+
     def test_start_run(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
         runner.start()
         runner.run_to_completion()
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(1), mock.call(1)])
+        self.assertEqual(2, self.mock_sleep.call_count)
+
     def test_start_run_wait_time(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(24).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(24).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
         runner.start()
         runner.run_to_completion(wait_time=24)
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(24), mock.call(24)])
+        self.assertEqual(2, self.mock_sleep.call_count)
+
     def test_run_progress(self):
         progress_count = []
 
@@ -630,21 +606,19 @@ class TaskTest(common.HeatTestCase):
             progress_count.append(None)
 
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(0).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         scheduler.TaskRunner(task)(progress_callback=progress)
         self.assertEqual(task.num_steps, len(progress_count))
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(0),
+                                          mock.call(1),
+                                          mock.call(1)])
+        self.assertEqual(3, self.mock_sleep.call_count)
+
     def test_start_run_progress(self):
         progress_count = []
 
@@ -652,22 +626,19 @@ class TaskTest(common.HeatTestCase):
             progress_count.append(None)
 
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
         runner.start()
         runner.run_to_completion(progress_callback=progress)
         self.assertEqual(task.num_steps - 1, len(progress_count))
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(1), mock.call(1)])
+        self.assertEqual(2, self.mock_sleep.call_count)
+
     def test_run_as_task_progress(self):
         progress_count = []
 
@@ -675,14 +646,7 @@ class TaskTest(common.HeatTestCase):
             progress_count.append(None)
 
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         tr = scheduler.TaskRunner(task)
         rt = tr.as_task(progress_callback=progress)
@@ -690,6 +654,11 @@ class TaskTest(common.HeatTestCase):
             pass
         self.assertEqual(task.num_steps, len(progress_count))
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+
     def test_run_progress_exception(self):
         class TestException(Exception):
             pass
@@ -702,20 +671,17 @@ class TaskTest(common.HeatTestCase):
             progress_count.append(None)
 
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(0).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         self.assertRaises(TestException, scheduler.TaskRunner(task),
                           progress_callback=progress)
         self.assertEqual(1, len(progress_count))
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2)])
+        self.assertEqual(2, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(0), mock.call(1)])
+        self.assertEqual(2, self.mock_sleep.call_count)
+
     def test_start_run_progress_exception(self):
         class TestException(Exception):
             pass
@@ -728,16 +694,7 @@ class TaskTest(common.HeatTestCase):
             progress_count.append(None)
 
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
         runner.start()
@@ -745,6 +702,12 @@ class TaskTest(common.HeatTestCase):
                           progress_callback=progress)
         self.assertEqual(1, len(progress_count))
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_has_calls([mock.call(1), mock.call(1)])
+        self.assertEqual(2, self.mock_sleep.call_count)
+
     def test_run_as_task_progress_exception(self):
         class TestException(Exception):
             pass
@@ -757,13 +720,7 @@ class TaskTest(common.HeatTestCase):
             progress_count.append(None)
 
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         tr = scheduler.TaskRunner(task)
         rt = tr.as_task(progress_callback=progress)
@@ -772,6 +729,10 @@ class TaskTest(common.HeatTestCase):
         self.assertRaises(TestException, next, rt)
         self.assertEqual(1, len(progress_count))
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2)])
+        self.assertEqual(2, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+
     def test_run_progress_exception_swallow(self):
         class TestException(Exception):
             pass
@@ -791,16 +752,12 @@ class TaskTest(common.HeatTestCase):
             except TestException:
                 yield
 
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        scheduler.TaskRunner._sleep(0).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
-
         scheduler.TaskRunner(task)(progress_callback=progress)
         self.assertEqual(2, len(progress_count))
 
+        self.mock_sleep.assert_has_calls([mock.call(0), mock.call(1)])
+        self.assertEqual(2, self.mock_sleep.call_count)
+
     def test_start_run_progress_exception_swallow(self):
         class TestException(Exception):
             pass
@@ -821,18 +778,14 @@ class TaskTest(common.HeatTestCase):
             except TestException:
                 yield
 
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-        scheduler.TaskRunner._sleep(1).AndReturn(None)
-
-        self.m.ReplayAll()
-
         runner = scheduler.TaskRunner(task)
         runner.start()
         runner.run_to_completion(progress_callback=progress)
         self.assertEqual(2, len(progress_count))
 
+        self.mock_sleep.assert_has_calls([mock.call(1), mock.call(1)])
+        self.assertEqual(2, self.mock_sleep.call_count)
+
     def test_run_as_task_progress_exception_swallow(self):
         class TestException(Exception):
             pass
@@ -859,60 +812,23 @@ class TaskTest(common.HeatTestCase):
         self.assertRaises(StopIteration, next, rt)
         self.assertEqual(2, len(progress_count))
 
-    def test_sleep(self):
-        sleep_time = 42
-        self.m.StubOutWithMock(eventlet, 'sleep')
-        eventlet.sleep(0).AndReturn(None)
-        eventlet.sleep(sleep_time).MultipleTimes().AndReturn(None)
-
-        self.m.ReplayAll()
-
-        runner = scheduler.TaskRunner(DummyTask())
-        runner(wait_time=sleep_time)
-
-    def test_sleep_zero(self):
-        self.m.StubOutWithMock(eventlet, 'sleep')
-        eventlet.sleep(0).MultipleTimes().AndReturn(None)
-
-        self.m.ReplayAll()
-
-        runner = scheduler.TaskRunner(DummyTask())
-        runner(wait_time=0)
-
-    def test_sleep_none(self):
-        self.m.StubOutWithMock(eventlet, 'sleep')
-        self.m.ReplayAll()
-
-        runner = scheduler.TaskRunner(DummyTask())
-        runner(wait_time=None)
-
     def test_args(self):
         args = ['foo', 'bar']
         kwargs = {'baz': 'quux', 'blarg': 'wibble'}
 
-        self.m.StubOutWithMock(DummyTask, '__call__')
-        task = DummyTask()
-
-        task(*args, **kwargs)
-
-        self.m.ReplayAll()
+        task = mock.Mock()
 
         runner = scheduler.TaskRunner(task, *args, **kwargs)
         runner(wait_time=None)
 
+        task.assert_called_with(*args, **kwargs)
+
     def test_non_callable(self):
         self.assertRaises(AssertionError, scheduler.TaskRunner, object())
 
     def test_stepping(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        task.do_step(3).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
         runner.start()
@@ -923,12 +839,14 @@ class TaskTest(common.HeatTestCase):
         self.assertTrue(runner.step())
         self.assertFalse(runner)
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3)])
+        self.assertEqual(3, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+
     def test_start_no_steps(self):
         task = DummyTask(0)
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
         runner.start()
@@ -936,14 +854,12 @@ class TaskTest(common.HeatTestCase):
         self.assertTrue(runner.done())
         self.assertTrue(runner.step())
 
+        task.do_step.assert_not_called()
+        self.mock_sleep.assert_not_called()
+
     def test_start_only(self):
         task = DummyTask()
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
 
@@ -951,6 +867,9 @@ class TaskTest(common.HeatTestCase):
         runner.start()
         self.assertTrue(runner.started())
 
+        task.do_step.assert_called_once_with(1)
+        self.mock_sleep.assert_not_called()
+
     def test_double_start(self):
         runner = scheduler.TaskRunner(DummyTask())
 
@@ -982,10 +901,7 @@ class TaskTest(common.HeatTestCase):
 
     def test_repeated_done(self):
         task = DummyTask(0)
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
 
@@ -993,6 +909,9 @@ class TaskTest(common.HeatTestCase):
         self.assertTrue(runner.step())
         self.assertTrue(runner.step())
 
+        task.do_step.assert_not_called()
+        self.mock_sleep.assert_not_called()
+
     def test_timeout(self):
         st = timeutils.wallclock()
 
@@ -1000,12 +919,8 @@ class TaskTest(common.HeatTestCase):
             while True:
                 yield
 
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.5, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
 
@@ -1013,6 +928,8 @@ class TaskTest(common.HeatTestCase):
         self.assertTrue(runner)
         self.assertRaises(scheduler.Timeout, runner.step)
 
+        self.assertEqual(3, timeutils.wallclock.call_count)
+
     def test_timeout_return(self):
         st = timeutils.wallclock()
 
@@ -1023,12 +940,8 @@ class TaskTest(common.HeatTestCase):
                 except scheduler.Timeout:
                     return
 
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.5, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
 
@@ -1037,6 +950,8 @@ class TaskTest(common.HeatTestCase):
         self.assertTrue(runner.step())
         self.assertFalse(runner)
 
+        self.assertEqual(3, timeutils.wallclock.call_count)
+
     def test_timeout_swallowed(self):
         st = timeutils.wallclock()
 
@@ -1048,12 +963,8 @@ class TaskTest(common.HeatTestCase):
                     yield
                     self.fail('Task still running')
 
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.5, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
 
@@ -1063,6 +974,8 @@ class TaskTest(common.HeatTestCase):
         self.assertFalse(runner)
         self.assertTrue(runner.step())
 
+        self.assertEqual(3, timeutils.wallclock.call_count)
+
     def test_as_task_timeout(self):
         st = timeutils.wallclock()
 
@@ -1070,12 +983,8 @@ class TaskTest(common.HeatTestCase):
             while True:
                 yield
 
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.5, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
 
@@ -1084,6 +993,8 @@ class TaskTest(common.HeatTestCase):
         self.assertTrue(runner)
         self.assertRaises(scheduler.Timeout, next, rt)
 
+        self.assertEqual(3, timeutils.wallclock.call_count)
+
     def test_as_task_timeout_shorter(self):
         st = timeutils.wallclock()
 
@@ -1091,14 +1002,9 @@ class TaskTest(common.HeatTestCase):
             while True:
                 yield
 
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        timeutils.wallclock().AndReturn(st + 0.7)
-        timeutils.wallclock().AndReturn(st + 1.6)
-        timeutils.wallclock().AndReturn(st + 2.6)
-
-        self.m.ReplayAll()
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.5, st + 0.7,
+                                      st + 1.6, st + 2.6])
 
         runner = scheduler.TaskRunner(task)
         runner.start(timeout=10)
@@ -1108,6 +1014,8 @@ class TaskTest(common.HeatTestCase):
         next(rt)
         self.assertRaises(scheduler.Timeout, next, rt)
 
+        self.assertEqual(5, timeutils.wallclock.call_count)
+
     def test_as_task_timeout_longer(self):
         st = timeutils.wallclock()
 
@@ -1115,13 +1023,8 @@ class TaskTest(common.HeatTestCase):
             while True:
                 yield
 
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        timeutils.wallclock().AndReturn(st + 0.6)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.5, st + 0.6, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
         runner.start(timeout=1)
@@ -1130,13 +1033,12 @@ class TaskTest(common.HeatTestCase):
         rt = runner.as_task(timeout=10)
         self.assertRaises(scheduler.Timeout, next, rt)
 
+        self.assertEqual(4, timeutils.wallclock.call_count)
+
     def test_cancel_not_started(self):
         task = DummyTask(1)
 
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
 
@@ -1146,15 +1048,13 @@ class TaskTest(common.HeatTestCase):
 
         self.assertTrue(runner.done())
 
+        task.do_step.assert_not_called()
+        self.mock_sleep.assert_not_called()
+
     def test_cancel_done(self):
         task = DummyTask(1)
 
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
 
@@ -1169,16 +1069,13 @@ class TaskTest(common.HeatTestCase):
         self.assertTrue(runner.done())
         self.assertTrue(runner.step())
 
+        task.do_step.assert_called_once_with(1)
+        self.mock_sleep.assert_not_called()
+
     def test_cancel(self):
         task = DummyTask(3)
 
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
 
@@ -1190,24 +1087,17 @@ class TaskTest(common.HeatTestCase):
         runner.cancel()
         self.assertTrue(runner.step())
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2)])
+        self.assertEqual(2, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+
     def test_cancel_grace_period(self):
         st = timeutils.wallclock()
         task = DummyTask(5)
 
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-
-        task.do_step(1).AndReturn(None)
-        task.do_step(2).AndReturn(None)
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        task.do_step(3).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 1.0)
-        task.do_step(4).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.5, st + 1.0, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
 
@@ -1221,27 +1111,20 @@ class TaskTest(common.HeatTestCase):
         self.assertFalse(runner.step())
         self.assertTrue(runner.step())
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3), mock.call(4)])
+        self.assertEqual(4, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+        self.assertEqual(4, timeutils.wallclock.call_count)
+
     def test_cancel_grace_period_before_timeout(self):
         st = timeutils.wallclock()
         task = DummyTask(5)
 
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.1)
-        task.do_step(1).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 0.2)
-        task.do_step(2).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 0.2)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        task.do_step(3).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 1.0)
-        task.do_step(4).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.1, st + 0.2, st + 0.2,
+                                      st + 0.5, st + 1.0, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
 
@@ -1255,27 +1138,20 @@ class TaskTest(common.HeatTestCase):
         self.assertFalse(runner.step())
         self.assertTrue(runner.step())
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3), mock.call(4)])
+        self.assertEqual(4, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+        self.assertEqual(7, timeutils.wallclock.call_count)
+
     def test_cancel_grace_period_after_timeout(self):
         st = timeutils.wallclock()
         task = DummyTask(5)
 
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-        self.m.StubOutWithMock(timeutils, 'wallclock')
-
-        timeutils.wallclock().AndReturn(st)
-        timeutils.wallclock().AndReturn(st + 0.1)
-        task.do_step(1).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 0.2)
-        task.do_step(2).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 0.2)
-        timeutils.wallclock().AndReturn(st + 0.5)
-        task.do_step(3).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 1.0)
-        task.do_step(4).AndReturn(None)
-        timeutils.wallclock().AndReturn(st + 1.5)
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
+        self.patchobject(timeutils, 'wallclock',
+                         side_effect=[st, st + 0.1, st + 0.2, st + 0.2,
+                                      st + 0.5, st + 1.0, st + 1.5])
 
         runner = scheduler.TaskRunner(task)
 
@@ -1289,13 +1165,16 @@ class TaskTest(common.HeatTestCase):
         self.assertFalse(runner.step())
         self.assertRaises(scheduler.Timeout, runner.step)
 
+        task.do_step.assert_has_calls([mock.call(1), mock.call(2),
+                                       mock.call(3), mock.call(4)])
+        self.assertEqual(4, task.do_step.call_count)
+        self.mock_sleep.assert_not_called()
+        self.assertEqual(7, timeutils.wallclock.call_count)
+
     def test_cancel_grace_period_not_started(self):
         task = DummyTask(1)
 
-        self.m.StubOutWithMock(task, 'do_step')
-        self.m.StubOutWithMock(scheduler.TaskRunner, '_sleep')
-
-        self.m.ReplayAll()
+        task.do_step = mock.Mock(return_value=None)
 
         runner = scheduler.TaskRunner(task)
 
@@ -1305,6 +1184,39 @@ class TaskTest(common.HeatTestCase):
 
         self.assertTrue(runner.done())
 
+        task.do_step.assert_not_called()
+        self.mock_sleep.assert_not_called()
+
+
+class TaskSleepTest(common.HeatTestCase):
+
+    def setUp(self):
+        super(TaskSleepTest, self).setUp()
+        scheduler.ENABLE_SLEEP = True
+        self.mock_sleep = self.patchobject(eventlet, 'sleep',
+                                           return_value=None)
+
+    def test_sleep(self):
+        sleep_time = 42
+
+        runner = scheduler.TaskRunner(DummyTask())
+        runner(wait_time=sleep_time)
+
+        self.mock_sleep.assert_any_call(0)
+        self.mock_sleep.assert_called_with(sleep_time)
+
+    def test_sleep_zero(self):
+        runner = scheduler.TaskRunner(DummyTask())
+        runner(wait_time=0)
+
+        self.mock_sleep.assert_called_with(0)
+
+    def test_sleep_none(self):
+        runner = scheduler.TaskRunner(DummyTask())
+        runner(wait_time=None)
+
+        self.mock_sleep.assert_not_called()
+
 
 class TimeoutTest(common.HeatTestCase):
     def test_compare(self):