def test_future_prop(self, mock_future, mock_client): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run rando function with a future handler future = client.submit(inc, 1) # Get the key from future handler and assign it to DaskJob key to keep track of this inc function djob.key = future.key # Use DaskJob future property to get back the inc function ret = djob.future # Check result mock_future.assert_called_with(key='test_key', client=mock_client_ret) self.assertEqual(mock_future(), ret)
def test_process_result_with_result_function(self, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): fake_key = 'sum_faef' mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function') djob.key = fake_key # call the function djob._process_results() # check the result mock_client.close.assert_called() mock_client.gather.assert_called_with(mock_future) mock_function.assert_called_with(mock_client.gather()) mock_client.set_metadata.assert_called_with(fake_key, False) self.assertEqual('', djob.key) mock_re_lock.assert_called()
def test_future_prop_exception(self, mock_future, mock_client, mock_log): mock_client_ret = mock.MagicMock() mock_client.return_value = mock_client_ret mock_client_ret.submit.return_value = mock.MagicMock(key='test_key') mock_future.side_effect = Exception('exception in creating future') # Create DaskJob djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler) # Get Scheduler Client from DaskJob using client property client = djob.client # Use this Client to run inc function with a future handler future = client.submit(inc, 1) # Get the key from future handler and assign it to DaskJob key to keep track of this inc function djob.key = future.key # Use DaskJob future property to get back the inc function ret = djob.future # Check result self.assertIsNone(ret) mock_log.exception.assert_called_with('Dask Future Init Error')
def test_process_result_with_result_function(self, mock_re_lock, mock_apl, mock_client, mock_future, mock_tfe): fake_key = 'sum_faef' mock_function_extractor = mock.MagicMock() mock_function = mock.MagicMock(return_value='foo') mock_function_extractor.valid = True mock_function_extractor.function = mock_function mock_tfe.return_value = mock_function_extractor mock_apl.return_value = True # Create DaskJob djob = DaskJob( name='test_dj', user=self.user, label='label', scheduler=self.scheduler, _process_results_function='test_function' ) djob.key = fake_key # call the function djob._process_results() # check the result mock_client.gather.assert_called_with(mock_future) mock_function.assert_called_with(mock_client.gather()) mock_client.set_metadata.assert_called_with(fake_key, False) self.assertEqual('', djob.key) mock_re_lock.assert_called()