diff --git a/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py b/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py index 88ff3c3d..c410bb47 100644 --- a/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py +++ b/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py @@ -62,9 +62,9 @@ class ShutDownScenarioPlugin(AbstractScenarioPlugin): return 1 def multiprocess_nodes(self, cloud_object_function, nodes, processes=0): + # pool object with number of element + pool = None try: - # pool object with number of element - if processes == 0: pool = ThreadPool(processes=len(nodes)) else: @@ -82,9 +82,12 @@ class ShutDownScenarioPlugin(AbstractScenarioPlugin): else: logging.info("pool type" + str(type(nodes))) pool.map(cloud_object_function, nodes) - pool.close() except Exception as e: - logging.info("Error on pool multiprocessing: " + str(e)) + logging.error("Error on pool multiprocessing: " + str(e)) + finally: + if pool: + pool.close() + pool.join() # Inject the cluster shut down scenario # krkn_lib diff --git a/tests/test_shut_down_scenario_plugin.py b/tests/test_shut_down_scenario_plugin.py index aad0ff82..44e2976d 100644 --- a/tests/test_shut_down_scenario_plugin.py +++ b/tests/test_shut_down_scenario_plugin.py @@ -409,7 +409,7 @@ class TestShutDownScenarioPlugin(unittest.TestCase): self.assertEqual(call_args[0], mock_cloud_function) mock_pool_instance.close.assert_called_once() - @patch('logging.info') + @patch('logging.error') @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.ThreadPool') def test_multiprocess_nodes_with_exception(self, mock_threadpool, mock_logging): """ @@ -425,6 +425,44 @@ class TestShutDownScenarioPlugin(unittest.TestCase): mock_logging.assert_called() logged_args, logged_kwargs = mock_logging.call_args self.assertIn("Error on pool multiprocessing", logged_args[0]) + + @patch('logging.error') + @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.ThreadPool') + def test_multiprocess_nodes_pool_close_on_exception(self, mock_threadpool, mock_logging): + """ + Test pool.close() and pool.join() are called even when map raises + """ + mock_pool_instance = Mock() + mock_threadpool.return_value = mock_pool_instance + mock_pool_instance.map.side_effect = Exception("map failed") + + nodes = ["node1", "node2"] + mock_cloud_function = Mock() + + self.plugin.multiprocess_nodes(mock_cloud_function, nodes, processes=0) + + mock_pool_instance.close.assert_called_once() + mock_pool_instance.join.assert_called_once() + + @patch('logging.error') + @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.ThreadPool') + def test_multiprocess_nodes_logs_error_on_exception(self, mock_threadpool, mock_logging): + """ + Test logging.error is used when an exception occurs + """ + mock_pool_instance = Mock() + mock_threadpool.return_value = mock_pool_instance + mock_pool_instance.map.side_effect = Exception("map failed") + + nodes = ["node1", "node2"] + mock_cloud_function = Mock() + + self.plugin.multiprocess_nodes(mock_cloud_function, nodes, processes=0) + + mock_logging.assert_called() + logged_message = mock_logging.call_args[0][0] + self.assertIn("Error on pool multiprocessing", logged_message) + @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.AWS') @patch('time.sleep') @patch('time.time')