def test_execute_query_succeeded_with_retries(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand, AlertQueryError
execute_query_mock = mocker.patch(
"superset.reports.commands.alert.AlertCommand._execute_query")
query_executed_count = 0
# Should match the value defined in superset_test_config.py
expected_max_retries = 3
def _mocked_execute_query() -> pd.DataFrame:
nonlocal query_executed_count
query_executed_count += 1
if query_executed_count < expected_max_retries:
raise AlertQueryError()
else:
return pd.DataFrame([{"sample_col": 0}])
execute_query_mock.side_effect = _mocked_execute_query
execute_query_mock.__name__ = "mocked_execute_query"
command = AlertCommand(report_schedule=mocker.Mock())
command.validate()
assert execute_query_mock.call_count == expected_max_retries
def test_execute_query_succeeded_no_retry(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand
execute_query_mock = mocker.patch(
"superset.reports.commands.alert.AlertCommand._execute_query",
side_effect=lambda: pd.DataFrame([{
"sample_col": 0
}]),
)
command = AlertCommand(report_schedule=mocker.Mock())
command.validate()
assert execute_query_mock.call_count == 1
def test_execute_query_failed_no_retry(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand, AlertQueryTimeout
execute_query_mock = mocker.patch(
"superset.reports.commands.alert.AlertCommand._execute_query")
def _mocked_execute_query() -> None:
raise AlertQueryTimeout
execute_query_mock.side_effect = _mocked_execute_query
execute_query_mock.__name__ = "mocked_execute_query"
command = AlertCommand(report_schedule=mocker.Mock())
try:
command.validate()
except AlertQueryTimeout:
pass
assert execute_query_mock.call_count == 1
def test_execute_query_failed_max_retries(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand, AlertQueryError
execute_query_mock = mocker.patch(
"superset.reports.commands.alert.AlertCommand._execute_query")
def _mocked_execute_query() -> None:
raise AlertQueryError
execute_query_mock.side_effect = _mocked_execute_query
execute_query_mock.__name__ = "mocked_execute_query"
command = AlertCommand(report_schedule=mocker.Mock())
try:
command.validate()
except AlertQueryError:
pass
# Should match the value defined in superset_test_config.py
assert execute_query_mock.call_count == 3