X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/8ed521f7fd1e48e1e415125745ed8c6627a62c91..0413abf93d7ee0cfd10222bf04aa88c4e7460303:/services/nodemanager/tests/test_failure.py

diff --git a/services/nodemanager/tests/test_failure.py b/services/nodemanager/tests/test_failure.py
index 285aa03c7d..ef4423dafa 100644
--- a/services/nodemanager/tests/test_failure.py
+++ b/services/nodemanager/tests/test_failure.py
@@ -1,4 +1,7 @@
 #!/usr/bin/env python
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: AGPL-3.0
 
 from __future__ import absolute_import, print_function
 
@@ -16,8 +19,8 @@ from . import testutil
 import arvnodeman.baseactor
 
 class BogusActor(arvnodeman.baseactor.BaseNodeManagerActor):
-    def __init__(self, e):
-        super(BogusActor, self).__init__()
+    def __init__(self, e, killfunc=None):
+        super(BogusActor, self).__init__(killfunc=killfunc)
         self.exp = e
 
     def doStuff(self):
@@ -32,24 +35,29 @@ class BogusActor(arvnodeman.baseactor.BaseNodeManagerActor):
 class ActorUnhandledExceptionTest(testutil.ActorTestMixin, unittest.TestCase):
     def test_fatal_error(self):
         for e in (MemoryError(), threading.ThreadError(), OSError(errno.ENOMEM, "")):
-            with mock.patch('os.kill') as kill_mock:
-                act = BogusActor.start(e).tell_proxy()
-                act.doStuff()
-                act.actor_ref.stop(block=True)
-                self.assertTrue(kill_mock.called)
-
-    @mock.patch('os.kill')
-    def test_nonfatal_error(self, kill_mock):
-        act = BogusActor.start(OSError(errno.ENOENT, "")).tell_proxy()
+            kill_mock = mock.Mock('os.kill')
+            bgact = BogusActor.start(e, killfunc=kill_mock)
+            act_thread = bgact.proxy().get_thread().get()
+            act = bgact.tell_proxy()
+            act.doStuff()
+            act.actor_ref.stop(block=True)
+            act_thread.join()
+            self.assertTrue(kill_mock.called)
+
+    def test_nonfatal_error(self):
+        kill_mock = mock.Mock('os.kill')
+        act = BogusActor.start(OSError(errno.ENOENT, ""), killfunc=kill_mock).tell_proxy()
         act.doStuff()
         act.actor_ref.stop(block=True)
         self.assertFalse(kill_mock.called)
 
 class WatchdogActorTest(testutil.ActorTestMixin, unittest.TestCase):
-    @mock.patch('os.kill')
-    def test_time_timout(self, kill_mock):
+
+    def test_time_timout(self):
+        kill_mock = mock.Mock('os.kill')
         act = BogusActor.start(OSError(errno.ENOENT, ""))
-        watch = arvnodeman.baseactor.WatchdogActor.start(1, act)
+        watch = arvnodeman.baseactor.WatchdogActor.start(1, act, killfunc=kill_mock)
+        time.sleep(1)
         watch.stop(block=True)
         act.stop(block=True)
         self.assertTrue(kill_mock.called)