def test_job_cloud_file(self):
# test that job_cloud_file() does not set job.status after yield
job = Job.objects.create()
job.status = Job.CLOUD_FILE_DOWNLOADED # an arbitrary one that's not PENDING (default) or SUCCESS
job.save()
with patch('utils.cloud_file.download_file') as download_file_mock:
# test when caller sets status - exposes bug where job_cloud_file() was setting to SUCCESS
with job_cloud_file(job.pk) as (job, cloud_file_fp):
job.status = Job.CLOUD_FILE_UPLOADED
job.save()
job.refresh_from_db()
self.assertEqual(Job.CLOUD_FILE_UPLOADED, job.status)
# test when __enter__() raises JobTimeoutException
job.status = Job.CLOUD_FILE_DOWNLOADED
job.save()
download_file_mock.side_effect = JobTimeoutException(
'download_file_mock Exception')
with self.assertRaises(JobTimeoutException):
with job_cloud_file(job.pk) as (job, cloud_file_fp):
pass
job.refresh_from_db()
self.assertEqual(Job.TIMEOUT, job.status)
# test when __enter__() raises BotoCoreError
job.status = Job.CLOUD_FILE_DOWNLOADED
job.save()
download_file_mock.side_effect = BotoCoreError()
with self.assertRaises(BotoCoreError):
with job_cloud_file(job.pk) as (job, cloud_file_fp):
pass
job.refresh_from_db()
self.assertEqual(Job.FAILED, job.status)
# test when __enter__() raises Exception
job.status = Job.CLOUD_FILE_DOWNLOADED
job.save()
download_file_mock.side_effect = Exception(
'download_file_mock Exception')
with self.assertRaises(Exception):
with job_cloud_file(job.pk) as (job, cloud_file_fp):
pass
job.refresh_from_db()
self.assertEqual(Job.FAILED, job.status)
def test__upload_forecast_worker_deletes_forecast(self):
# verifies that _upload_forecast_worker() deletes the (presumably empty) Forecast that's passed to it by
# upload functions if the file is invalid. here we mock load_predictions_from_json_io_dict() to throw the two
# exceptions that cause deletes: JobTimeoutException and Exception
_, _, po_user, _, _, _, _, _ = get_or_create_super_po_mo_users(
is_create_super=True)
project, time_zero, forecast_model, forecast = _make_docs_project(
po_user)
forecast.issued_at -= datetime.timedelta(
days=1) # older version avoids unique constraint errors
forecast.save()
for exception, exp_job_status in [
(Exception('load_preds_mock Exception'), Job.FAILED),
(JobTimeoutException('load_preds_mock JobTimeoutException'),
Job.TIMEOUT)
]:
with patch('forecast_app.models.job.job_cloud_file') as job_cloud_file_mock, \
patch('utils.forecast.load_predictions_from_json_io_dict') as load_preds_mock, \
patch('utils.forecast.cache_forecast_metadata') as cache_metatdata_mock, \
open('forecast_app/tests/predictions/docs-predictions.json') as cloud_file_fp:
load_preds_mock.side_effect = exception
forecast2 = Forecast.objects.create(
forecast_model=forecast_model, time_zero=time_zero)
job = Job.objects.create()
job.input_json = {
'forecast_pk': forecast2.pk,
'filename': 'a name!'
}
job.save()
job_cloud_file_mock.return_value.__enter__.return_value = (
job, cloud_file_fp)
try:
_upload_forecast_worker(job.pk)
except JobTimeoutException as jte:
pass # expected re-raise of this exception
job.refresh_from_db()
self.assertEqual(exp_job_status, job.status)
self.assertIsNone(
Forecast.objects.filter(
id=forecast2.id).first()) # deleted
def test_create_tasks_throws_timeout_exception(self, importer, send_mail):
importer.create_tasks.side_effect = JobTimeoutException()
uploader_name = 'Cersei Lannister'
project = ProjectFactory.create()
form_data = {'type': 'csv', 'csv_url': 'http://google.es'}
subject = 'Your import task has timed out'
with patch.dict(self.flask_app.config, {'BRAND': 'GOT'}):
body = '\n'.join(
['Hello,\n',
'Import task to your project {} by {} failed because the file was too large.',
'It was able to process approximately {} tasks.',
'Please break up your task upload into smaller CSV files.',
'Thank you,\n',
u'The {} team.']).format(project.name, uploader_name,
0, self.flask_app.config['BRAND'])
email_data = dict(recipients=[project.owner.email_addr],
subject=subject, body=body)
assert_raises(JobTimeoutException, import_tasks, project.id, uploader_name, **form_data)
send_mail.assert_called_once_with(email_data)
def setup_death_penalty(self):
exception = JobTimeoutException(
'Gevent Job exceeded maximum timeout value (%d seconds).' %
self._timeout)
self.gevent_timeout = gevent.Timeout(self._timeout, exception)
self.gevent_timeout.start()