def test_process_result_serialize_exception(self, mock_log, mock_re_lock,
mock_apl, mock_client,
mock_future, mock_tfe):
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')
# NOTE: To mock the "result" property, we must mock it on the type object, not the instance.
# Unfortunately, this will persist for any test in the same test case that runs after this test.
# That's why this test is pulled out in a separate test case (the other tests on "result" won't work
# after this one runs).
type(djob).result = mock.PropertyMock(
side_effect=[Exception, 'foo', 'foo'])
djob._process_results()
# check the result
mock_client.gather.assert_called_with(mock_future)
mock_function.assert_called_with(mock_client.gather())
mock_log.exception.assert_called_with('Results Serialization Error')
mock_re_lock.assert_called()
self.assertEqual('ERR', djob._status)
def test_process_result_with_result_function_with_exception(
self, mock_re_lock, mock_apl, _, mock_client, mock_tfe, mock_log,
mock_save):
mock_function_extractor = mock.MagicMock()
mock_function = mock.MagicMock()
mock_function.side_effect = Exception
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')
# call the function
djob._process_results()
# check the result
mock_log.exception.assert_called_with('Process Results Function Error')
self.assertEqual('ERR', djob._status)
mock_save.assert_called()
mock_re_lock.assert_called()
def test_process_result_serialize_exception(self, mock_log, mock_re_lock, mock_apl, mock_client, mock_future,
mock_tfe):
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'
)
# NOTE: To mock the "result" property, we must mock it on the type object, not the instance.
# Unfortunately, this will persist for any test in the same test case that runs after this test.
# That's why this test is pulled out in a separate test case (the other tests on "result" won't work
# after this one runs).
type(djob).result = mock.PropertyMock(side_effect=[Exception, 'foo', 'foo'])
djob._process_results()
# check the result
mock_client.gather.assert_called_with(mock_future)
mock_function.assert_called_with(mock_client.gather())
mock_log.exception.assert_called_with('Results Serialization Error')
mock_re_lock.assert_called()
self.assertEqual('ERR', djob._status)
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_process_result_with_client_gather_exception(
self, mock_logger, mock_re_lock, mock_apl, mock_client,
mock_future, mock_tfe):
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
gather_exception = Exception('Fake exception')
mock_client.gather.side_effect = gather_exception
# Create DaskJob
djob = DaskJob(name='test_dj',
user=self.user,
label='label',
scheduler=self.scheduler,
_process_results_function='test_function')
# call the function
djob._process_results()
# check the result
mock_client.gather.assert_called_with(mock_future)
mock_logger.warning.assert_called()
mock_function.assert_called_with(gather_exception)
mock_re_lock.assert_called()
def test_process_result_with_client_gather_exception(self, mock_logger, mock_re_lock, mock_apl, mock_client,
mock_future, mock_tfe):
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
gather_exception = Exception('Fake exception')
mock_client.gather.side_effect = gather_exception
# Create DaskJob
djob = DaskJob(
name='test_dj',
user=self.user,
label='label',
scheduler=self.scheduler,
_process_results_function='test_function'
)
# call the function
djob._process_results()
# check the result
mock_client.gather.assert_called_with(mock_future)
mock_logger.warning.assert_called()
mock_function.assert_called_with(gather_exception)
mock_re_lock.assert_called()
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()
def test_process_result_forget(self, _):
# Create DaskJob
djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, forget=True)
# call the function
ret = djob._process_results()
# check the result
self.assertIsNone(ret)
def test_process_result_no_future(self, mock_apl, _):
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')
# call the function
self.assertIsNone(djob._process_results())
def test_process_result_no_future(self, mock_apl, _):
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')
# call the function
self.assertIsNone(djob._process_results())
def test_process_result_forget(self, _, mock_client):
# Create DaskJob
djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler, forget=True)
# call the function
ret = djob._process_results()
# check the result
mock_client.close.assert_called()
self.assertIsNone(ret)
def test_process_result_with_failed_lock(self, mock_re_lock, mock_apl):
mock_apl.return_value = False
# Create DaskJob
djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler,
_process_results_function='test_function')
# call the function
self.assertIsNone(djob._process_results())
# check the result
mock_re_lock.assert_not_called()
def test_process_result_with_failed_lock(self, mock_re_lock, mock_apl):
mock_apl.return_value = False
# Create DaskJob
djob = DaskJob(name='test_dj', user=self.user, label='label', scheduler=self.scheduler,
_process_results_function='test_function')
# call the function
self.assertIsNone(djob._process_results())
# check the result
mock_re_lock.assert_not_called()
def test_process_result_with_result_function_with_exception(self, mock_re_lock, mock_apl, _, mock_client,
mock_tfe, mock_log, mock_save):
mock_function_extractor = mock.MagicMock()
mock_function = mock.MagicMock()
mock_function.side_effect = Exception
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')
# call the function
djob._process_results()
# check the result
mock_log.exception.assert_called_with('Process Results Function Error')
self.assertEquals('ERR', djob._status)
mock_save.assert_called()
mock_re_lock.assert_called()
