Files
superset2/tests/unit_tests/test_viz_get_df_payload.py

112 lines
3.8 KiB
Python

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from typing import Any
from unittest.mock import patch
import pytest
from superset import viz
from superset.common.db_query_status import QueryStatus
from superset.connectors.sqla.models import SqlaTable
from superset.errors import SupersetErrorType
from superset.exceptions import (
OAuth2RedirectError,
QueryObjectValidationError,
)
from superset.models.core import Database
QUERY_OBJ: dict[str, Any] = {"row_limit": 100, "from_dttm": None, "to_dttm": None}
def _viz() -> viz.BaseViz:
database = Database(database_name="d", sqlalchemy_uri="sqlite://")
datasource = SqlaTable(
table_name="t",
columns=[],
metrics=[],
main_dttm_col=None,
database=database,
)
# ``force=True`` skips the data cache lookup branch so ``get_df`` is always
# invoked, which is what we want to assert error-handling against.
return viz.BaseViz(
datasource=datasource,
form_data={"viz_type": "table"},
force=True,
)
def test_get_df_payload_propagates_oauth2_redirect_error() -> None:
"""
OAuth2RedirectError (a SupersetErrorException) must propagate out of
``get_df_payload`` so the global Flask error handler can serialize it.
"""
obj = _viz()
oauth_exc = OAuth2RedirectError(
url="https://accounts.example.com/o/oauth2/v2/auth?...",
tab_id="tab-123",
redirect_uri="https://superset.example.com/oauth2/redirect",
)
with patch.object(viz.BaseViz, "get_df", side_effect=oauth_exc):
with pytest.raises(OAuth2RedirectError) as exc_info:
obj.get_df_payload(QUERY_OBJ)
assert exc_info.value.error.error_type == SupersetErrorType.OAUTH2_REDIRECT
assert exc_info.value.error.extra == {
"url": "https://accounts.example.com/o/oauth2/v2/auth?...",
"tab_id": "tab-123",
"redirect_uri": "https://superset.example.com/oauth2/redirect",
}
def test_get_df_payload_captures_generic_exception_as_viz_get_df_error() -> None:
"""
Non-Superset exception raised by ``get_df`` are downgraded to a
``VIZ_GET_DF_ERROR`` entry on ``self.errors``.
"""
obj = _viz()
with patch.object(viz.BaseViz, "get_df", side_effect=RuntimeError("boom")):
payload = obj.get_df_payload(QUERY_OBJ)
assert obj.status == QueryStatus.FAILED
assert payload["status"] == QueryStatus.FAILED
assert len(obj.errors) == 1
assert obj.errors[0]["error_type"] == SupersetErrorType.VIZ_GET_DF_ERROR
assert obj.errors[0]["message"] == "boom"
def test_get_df_payload_captures_query_object_validation_error() -> None:
"""
``QueryObjectValidationError`` is reported as ``VIZ_GET_DF_ERROR``.
"""
obj = _viz()
with patch.object(
viz.BaseViz,
"get_df",
side_effect=QueryObjectValidationError("bad query"),
):
payload = obj.get_df_payload(QUERY_OBJ)
assert obj.status == QueryStatus.FAILED
assert payload["status"] == QueryStatus.FAILED
assert len(obj.errors) == 1
assert obj.errors[0]["error_type"] == SupersetErrorType.VIZ_GET_DF_ERROR
assert obj.errors[0]["message"] == "bad query"