def test_setup_multithreading(self):
with patch("maggma.runner.Thread") as mock_thread:
proc = MultiprocProcessor([self.builder], num_workers=3)
proc.builder = proc.builders[0]
proc.setup_multithreading()
mock_thread.assert_called()
def test_put_tasks(self):
with patch("maggma.runner.ProcessPoolExecutor") as mock_executor:
mock_exec_obj = mock_executor()
proc = MultiprocProcessor([self.builder], num_workers=3)
proc.builder = MagicMock()
proc.task_count = MagicMock()
cursor = [True, True, True, False]
proc.get_pbar = cursor
proc.put_tasks()
proc.task_count.acquire.assert_called()
def test_clean_up_data(self):
proc = MultiprocProcessor([self.builder], num_workers=3)
proc.data = MagicMock()
proc.update_data_condition = MagicMock()
proc.builder = MagicMock()
proc.update_targets_thread = MagicMock()
proc.clean_up_data()
proc.update_data_condition.notify_all.assert_called()
proc.update_targets_thread.join.assert_called()
def test_update_data_callback(self):
proc = MultiprocProcessor([self.builder], num_workers=3)
future = MagicMock()
proc.data = MagicMock()
proc.task_count = MagicMock()
proc.update_data_condition = MagicMock()
proc.process_pbar = MagicMock()
proc.update_data_callback(future)
future.result.assert_called()
proc.update_data_condition.notify_all.assert_called()
proc.task_count.release.assert_called()
def test_update_targets(self):
proc = MultiprocProcessor([self.builder], num_workers=3)
proc.builder = self.builder
proc.update_data_condition = MagicMock()
proc.data = MagicMock()
proc.update_pbar = MagicMock()
proc.run_update_targets = MagicMock()
proc.run_update_targets.__bool__.side_effect = [
True, True, True, False
]
proc.update_targets()
proc.run_update_targets.__bool__.assert_called()
proc.data.clear.assert_called()
proc.update_data_condition.wait_for.assert_called()
proc.builder.update_targets.assert_called()
def test_init(self):
proc = MultiprocProcessor([], 3)
self.assertEqual(proc.num_workers, 3)