From 009c43afacd5171f25ec53087662428b2d7690f1 Mon Sep 17 00:00:00 2001 From: Lily Kuang Date: Wed, 7 Apr 2021 16:15:21 -0700 Subject: [PATCH] update test --- superset/reports/commands/alert.py | 2 +- superset/reports/commands/execute.py | 2 +- superset/sql_lab.py | 2 +- superset/tasks/async_queries.py | 8 +++-- superset/tasks/scheduler.py | 2 +- tests/reports/commands_tests.py | 28 +++++++++++---- tests/sqllab_tests.py | 30 +++++++++++++++- tests/tasks/async_queries_tests.py | 51 ++++++++++++++++++++++++++++ tests/thumbnails_tests.py | 24 ++++++------- 9 files changed, 124 insertions(+), 25 deletions(-) diff --git a/superset/reports/commands/alert.py b/superset/reports/commands/alert.py index f7991bd3d3d36..5dd9797055c1f 100644 --- a/superset/reports/commands/alert.py +++ b/superset/reports/commands/alert.py @@ -156,7 +156,7 @@ def _execute_query(self) -> pd.DataFrame: ) return df except SoftTimeLimitExceeded as ex: - logger.error("A timeout occurred while executing the alert query: %s", ex) + logger.warning("A timeout occurred while executing the alert query: %s", ex) raise AlertQueryTimeout() except Exception as ex: raise AlertQueryError(message=str(ex)) diff --git a/superset/reports/commands/execute.py b/superset/reports/commands/execute.py index 5a68b02ec0bed..93037357895d6 100644 --- a/superset/reports/commands/execute.py +++ b/superset/reports/commands/execute.py @@ -181,7 +181,7 @@ def _get_screenshot(self) -> bytes: try: image_data = screenshot.get_screenshot(user=user) except SoftTimeLimitExceeded: - logger.error("A timeout occurred while taking a screenshot.") + logger.warning("A timeout occurred while taking a screenshot.") raise ReportScheduleScreenshotTimeout() except Exception as ex: raise ReportScheduleScreenshotFailedError( diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 6b37568130be8..234b1dd110e62 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -160,7 +160,7 @@ def get_sql_results( # pylint: disable=too-many-arguments log_params=log_params, ) except SoftTimeLimitExceeded as ex: - logger.error("Query %d: Time limit exceeded", query_id) + logger.warning("Query %d: Time limit exceeded", query_id) logger.debug("Query %d: %s", query_id, ex) raise SqlLabTimeoutException( _( diff --git a/superset/tasks/async_queries.py b/superset/tasks/async_queries.py index 13ca00a9e3f36..f008bc1f17d52 100644 --- a/superset/tasks/async_queries.py +++ b/superset/tasks/async_queries.py @@ -62,7 +62,9 @@ def load_chart_data_into_cache( job_metadata, async_query_manager.STATUS_DONE, result_url=result_url, ) except SoftTimeLimitExceeded as exc: - logger.error("A timeout occurred while loading chart data, error: %s", exc) + logger.warning( + "A timeout occurred while loading chart data, error: %s", exc + ) raise exc except Exception as exc: # TODO: QueryContext should support SIP-40 style errors @@ -109,7 +111,9 @@ def load_explore_json_into_cache( # pylint: disable=too-many-locals job_metadata, async_query_manager.STATUS_DONE, result_url=result_url, ) except SoftTimeLimitExceeded as ex: - logger.error("A timeout occurred while loading explore json, error: %s", ex) + logger.warning( + "A timeout occurred while loading explore json, error: %s", ex + ) raise ex except Exception as exc: if isinstance(exc, SupersetVizException): diff --git a/superset/tasks/scheduler.py b/superset/tasks/scheduler.py index 03fe72380505f..86bc9ca180271 100644 --- a/superset/tasks/scheduler.py +++ b/superset/tasks/scheduler.py @@ -93,6 +93,6 @@ def prune_log() -> None: try: AsyncPruneReportScheduleLogCommand().run() except SoftTimeLimitExceeded as ex: - logger.error("A timeout occurred while pruning report schedule logs: %s", ex) + logger.warning("A timeout occurred while pruning report schedule logs: %s", ex) except CommandException as ex: logger.error("An exception occurred while pruning report schedule logs: %s", ex) diff --git a/tests/reports/commands_tests.py b/tests/reports/commands_tests.py index 7dfb70c48cee5..162b9a87f3c0c 100644 --- a/tests/reports/commands_tests.py +++ b/tests/reports/commands_tests.py @@ -47,11 +47,13 @@ ReportScheduleNotFoundError, ReportScheduleNotificationError, ReportSchedulePreviousWorkingError, + ReportSchedulePruneLogError, ReportScheduleScreenshotFailedError, ReportScheduleScreenshotTimeout, ReportScheduleWorkingTimeoutError, ) from superset.reports.commands.execute import AsyncExecuteReportScheduleCommand +from superset.reports.commands.log_prune import AsyncPruneReportScheduleLogCommand from superset.utils.core import get_example_database from tests.fixtures.birth_names_dashboard import load_birth_names_dashboard_with_slices from tests.fixtures.world_bank_dashboard import ( @@ -186,7 +188,7 @@ def create_test_table_context(database: Database): database.get_sqla_engine().execute("DROP TABLE test_table") -@pytest.yield_fixture() +@pytest.fixture() def create_report_email_chart(): with app.app_context(): chart = db.session.query(Slice).first() @@ -198,7 +200,7 @@ def create_report_email_chart(): cleanup_report_schedule(report_schedule) -@pytest.yield_fixture() +@pytest.fixture() def create_report_email_dashboard(): with app.app_context(): dashboard = db.session.query(Dashboard).first() @@ -210,7 +212,7 @@ def create_report_email_dashboard(): cleanup_report_schedule(report_schedule) -@pytest.yield_fixture() +@pytest.fixture() def create_report_slack_chart(): with app.app_context(): chart = db.session.query(Slice).first() @@ -222,7 +224,7 @@ def create_report_slack_chart(): cleanup_report_schedule(report_schedule) -@pytest.yield_fixture() +@pytest.fixture() def create_report_slack_chart_working(): with app.app_context(): chart = db.session.query(Slice).first() @@ -248,7 +250,7 @@ def create_report_slack_chart_working(): cleanup_report_schedule(report_schedule) -@pytest.yield_fixture() +@pytest.fixture() def create_alert_slack_chart_success(): with app.app_context(): chart = db.session.query(Slice).first() @@ -274,7 +276,7 @@ def create_alert_slack_chart_success(): cleanup_report_schedule(report_schedule) -@pytest.yield_fixture() +@pytest.fixture() def create_alert_slack_chart_grace(): with app.app_context(): chart = db.session.query(Slice).first() @@ -1109,3 +1111,17 @@ def test_grace_period_error_flap( assert ( get_notification_error_sent_count(create_invalid_sql_alert_email_chart) == 2 ) + + +@pytest.mark.usefixtures( + "load_birth_names_dashboard_with_slices", "create_report_email_dashboard" +) +@patch("superset.reports.dao.ReportScheduleDAO.bulk_delete_logs") +def test_prune_log_soft_time_out(bulk_delete_logs, create_report_email_dashboard): + from celery.exceptions import SoftTimeLimitExceeded + from datetime import datetime, timedelta + + bulk_delete_logs.side_effect = SoftTimeLimitExceeded() + with pytest.raises(SoftTimeLimitExceeded) as excinfo: + AsyncPruneReportScheduleLogCommand().run() + assert str(excinfo.value) == "SoftTimeLimitExceeded()" diff --git a/tests/sqllab_tests.py b/tests/sqllab_tests.py index ee0fac541a346..afed7a678a530 100644 --- a/tests/sqllab_tests.py +++ b/tests/sqllab_tests.py @@ -33,7 +33,12 @@ from superset.models.core import Database from superset.models.sql_lab import Query, SavedQuery from superset.result_set import SupersetResultSet -from superset.sql_lab import execute_sql_statements, SqlLabException +from superset.sql_lab import ( + execute_sql_statements, + get_sql_results, + SqlLabException, + SqlLabTimeoutException, +) from superset.sql_parse import CtasMethod from superset.utils.core import ( datetime_to_epoch, @@ -793,3 +798,26 @@ def test_execute_sql_statements_ctas( "sure your query has only a SELECT statement. Then, " "try running your query again." ) + + @mock.patch("superset.sql_lab.get_query") + @mock.patch("superset.sql_lab.execute_sql_statement") + def test_get_sql_results_soft_time_limit( + self, mock_execute_sql_statement, mock_get_query + ): + from celery.exceptions import SoftTimeLimitExceeded + + sql = """ + -- comment + SET @value = 42; + SELECT @value AS foo; + -- comment + """ + mock_get_query.side_effect = SoftTimeLimitExceeded() + with pytest.raises(SqlLabTimeoutException) as excinfo: + get_sql_results( + 1, sql, return_results=True, store_results=False, + ) + assert ( + str(excinfo.value) + == "SQL Lab timeout. This environment's policy is to kill queries after 21600 seconds." + ) diff --git a/tests/tasks/async_queries_tests.py b/tests/tasks/async_queries_tests.py index cd4f0c0ce7669..cca58d8e2a0af 100644 --- a/tests/tasks/async_queries_tests.py +++ b/tests/tasks/async_queries_tests.py @@ -20,6 +20,7 @@ from uuid import uuid4 import pytest +from celery.exceptions import SoftTimeLimitExceeded from superset import db from superset.charts.commands.data import ChartDataCommand @@ -94,6 +95,31 @@ def test_load_chart_data_into_cache_error(self, mock_update_job, mock_run_comman errors = [{"message": "Error: foo"}] mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors) + @mock.patch.object(ChartDataCommand, "run") + @mock.patch.object(async_query_manager, "update_job") + def test_soft_timeout_load_chart_data_into_cache( + self, mock_update_job, mock_run_command + ): + async_query_manager.init_app(app) + user = security_manager.find_user("gamma") + form_data = {} + job_metadata = { + "channel_id": str(uuid4()), + "job_id": str(uuid4()), + "user_id": user.id, + "status": "pending", + "errors": [], + } + errors = ["A timeout occurred while loading chart data"] + + with pytest.raises(SoftTimeLimitExceeded): + with mock.patch.object( + async_queries, "ensure_user_is_set", + ) as ensure_user_is_set: + ensure_user_is_set.side_effect = SoftTimeLimitExceeded() + load_chart_data_into_cache(job_metadata, form_data) + ensure_user_is_set.assert_called_once_with(user.id, "error", errors=errors) + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") @mock.patch.object(async_query_manager, "update_job") def test_load_explore_json_into_cache(self, mock_update_job): @@ -151,3 +177,28 @@ def test_load_explore_json_into_cache_error(self, mock_update_job): errors = ["The dataset associated with this chart no longer exists"] mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors) + + @mock.patch.object(ChartDataCommand, "run") + @mock.patch.object(async_query_manager, "update_job") + def test_soft_timeout_load_explore_json_into_cache( + self, mock_update_job, mock_run_command + ): + async_query_manager.init_app(app) + user = security_manager.find_user("gamma") + form_data = {} + job_metadata = { + "channel_id": str(uuid4()), + "job_id": str(uuid4()), + "user_id": user.id, + "status": "pending", + "errors": [], + } + errors = ["A timeout occurred while loading explore json, error"] + + with pytest.raises(SoftTimeLimitExceeded): + with mock.patch.object( + async_queries, "ensure_user_is_set", + ) as ensure_user_is_set: + ensure_user_is_set.side_effect = SoftTimeLimitExceeded() + load_explore_json_into_cache(job_metadata, form_data) + ensure_user_is_set.assert_called_once_with(user.id, "error", errors=errors) diff --git a/tests/thumbnails_tests.py b/tests/thumbnails_tests.py index fb1fd689aa909..5879f4a8b6437 100644 --- a/tests/thumbnails_tests.py +++ b/tests/thumbnails_tests.py @@ -48,7 +48,7 @@ def url_open_auth(self, username: str, url: str): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_async_dashboard_screenshot(self): """ - Thumbnails: Simple get async dashboard screenshot + Thumbnails: Simple get async dashboard screenshot """ dashboard = db.session.query(Dashboard).all()[0] with patch("superset.dashboards.api.DashboardRestApi.get") as mock_get: @@ -65,7 +65,7 @@ class TestThumbnails(SupersetTestCase): def test_dashboard_thumbnail_disabled(self): """ - Thumbnails: Dashboard thumbnail disabled + Thumbnails: Dashboard thumbnail disabled """ if is_feature_enabled("THUMBNAILS"): return @@ -77,7 +77,7 @@ def test_dashboard_thumbnail_disabled(self): def test_chart_thumbnail_disabled(self): """ - Thumbnails: Chart thumbnail disabled + Thumbnails: Chart thumbnail disabled """ if is_feature_enabled("THUMBNAILS"): return @@ -90,7 +90,7 @@ def test_chart_thumbnail_disabled(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_async_dashboard_screenshot(self): """ - Thumbnails: Simple get async dashboard screenshot + Thumbnails: Simple get async dashboard screenshot """ dashboard = db.session.query(Dashboard).all()[0] self.login(username="admin") @@ -105,7 +105,7 @@ def test_get_async_dashboard_screenshot(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_async_dashboard_notfound(self): """ - Thumbnails: Simple get async dashboard not found + Thumbnails: Simple get async dashboard not found """ max_id = db.session.query(func.max(Dashboard.id)).scalar() self.login(username="admin") @@ -116,7 +116,7 @@ def test_get_async_dashboard_notfound(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_async_dashboard_not_allowed(self): """ - Thumbnails: Simple get async dashboard not allowed + Thumbnails: Simple get async dashboard not allowed """ dashboard = db.session.query(Dashboard).all()[0] self.login(username="gamma") @@ -127,7 +127,7 @@ def test_get_async_dashboard_not_allowed(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_async_chart_screenshot(self): """ - Thumbnails: Simple get async chart screenshot + Thumbnails: Simple get async chart screenshot """ chart = db.session.query(Slice).all()[0] self.login(username="admin") @@ -142,7 +142,7 @@ def test_get_async_chart_screenshot(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_async_chart_notfound(self): """ - Thumbnails: Simple get async chart not found + Thumbnails: Simple get async chart not found """ max_id = db.session.query(func.max(Slice.id)).scalar() self.login(username="admin") @@ -153,7 +153,7 @@ def test_get_async_chart_notfound(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_cached_chart_wrong_digest(self): """ - Thumbnails: Simple get chart with wrong digest + Thumbnails: Simple get chart with wrong digest """ chart = db.session.query(Slice).all()[0] chart_url = get_url_path("Superset.slice", slice_id=chart.id, standalone="true") @@ -169,7 +169,7 @@ def test_get_cached_chart_wrong_digest(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_cached_dashboard_screenshot(self): """ - Thumbnails: Simple get cached dashboard screenshot + Thumbnails: Simple get cached dashboard screenshot """ dashboard = db.session.query(Dashboard).all()[0] dashboard_url = get_url_path("Superset.dashboard", dashboard_id=dashboard.id) @@ -185,7 +185,7 @@ def test_get_cached_dashboard_screenshot(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_cached_chart_screenshot(self): """ - Thumbnails: Simple get cached chart screenshot + Thumbnails: Simple get cached chart screenshot """ chart = db.session.query(Slice).all()[0] chart_url = get_url_path("Superset.slice", slice_id=chart.id, standalone="true") @@ -201,7 +201,7 @@ def test_get_cached_chart_screenshot(self): @skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature") def test_get_cached_dashboard_wrong_digest(self): """ - Thumbnails: Simple get dashboard with wrong digest + Thumbnails: Simple get dashboard with wrong digest """ dashboard = db.session.query(Dashboard).all()[0] dashboard_url = get_url_path("Superset.dashboard", dashboard_id=dashboard.id)