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_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.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')
