[tasks] Mock celery task to enable tests without celery
@@ -50,8 +50,11 @@ class TestRunAsyncQueryTask():
with patch('notebook.tasks.get_api') as get_api:
with patch('notebook.tasks.DataAdapter') as DataAdapter:
with patch('notebook.tasks.export_csvxls.create_generator') as create_generator:
+ with patch('notebook.tasks.caches') as caches:
- DataAdapter.return_value = MagicMock(row_counter=2)
+ content_generator = MagicMock(row_counter=2)
+ content_generator.__iter__.return_value = [('col1', [[1]]), ('col2', [[2]])]
+ DataAdapter.return_value = content_generator
get_api.return_value = Mock(
check_status=Mock(return_value={'status': 0})
@@ -73,6 +76,8 @@ class TestRunAsyncQueryTask():
def test_close_statement(self):
with patch('notebook.tasks._get_request') as _get_request:
with patch('notebook.tasks.download_to_file') as download_to_file:
+ with patch('notebook.tasks.close_statement_async') as close_statement_async:
+ with patch('notebook.tasks.caches') as caches:
download_to_file.AsyncResult.return_value = Mock(
state=states.SUCCESS
