QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-18 23:55:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: set Snowflake user agent (#22432)

Browse Source
This commit is contained in:
Beto Dealmeida
2022-12-15 17:08:34 -08:00
committed by GitHub
parent 4f9c2c8165
commit bdeedaaf80
5 changed files with 107 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
tests/unit_tests/db_engine_specs/test_databricks.py
View File

@@ -18,8 +18,9 @@
import json
from pytest_mock import MockerFixture
from superset.utils.core import GenericDataType
from tests.integration_tests.db_engine_specs.base_tests import assert_generic_types
def test_get_parameters_from_uri() -> None:
@@ -110,6 +111,7 @@ def test_generic_type() -> None:
assert that generic types match
"""
from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
from tests.integration_tests.db_engine_specs.base_tests import assert_generic_types
type_expectations = (
# Numeric
@@ -133,3 +135,43 @@ def test_generic_type() -> None:
("BOOLEAN", GenericDataType.BOOLEAN),
)
assert_generic_types(DatabricksNativeEngineSpec, type_expectations)
def test_get_extra_params(mocker: MockerFixture) -> None:
"""
Test the ``get_extra_params`` method.
"""
from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
database = mocker.MagicMock()
database.extra = {}
assert DatabricksNativeEngineSpec.get_extra_params(database) == {
"engine_params": {
"connect_args": {
"http_headers": [("User-Agent", "Apache Superset")],
"_user_agent_entry": "Apache Superset",
}
}
}
database.extra = json.dumps(
{
"engine_params": {
"connect_args": {
"http_headers": [("User-Agent", "Custom user agent")],
"_user_agent_entry": "Custom user agent",
"foo": "bar",
}
}
}
)
assert DatabricksNativeEngineSpec.get_extra_params(database) == {
"engine_params": {
"connect_args": {
"http_headers": [["User-Agent", "Custom user agent"]],
"_user_agent_entry": "Custom user agent",
"foo": "bar",
}
}
}

31
tests/unit_tests/db_engine_specs/test_snowflake.py
View File

@@ -14,11 +14,15 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=import-outside-toplevel
import json
from datetime import datetime
from unittest import mock
import pytest
from pytest_mock import MockerFixture
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from tests.unit_tests.fixtures.common import dttm
@@ -122,3 +126,30 @@ def test_cancel_query_failed(engine_mock: mock.Mock) -> None:
query = Query()
cursor_mock = engine_mock.raiseError.side_effect = Exception()
assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, "123") is False
def test_get_extra_params(mocker: MockerFixture) -> None:
"""
Test the ``get_extra_params`` method.
"""
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
database = mocker.MagicMock()
database.extra = {}
assert SnowflakeEngineSpec.get_extra_params(database) == {
"engine_params": {"connect_args": {"application": "Apache Superset"}}
}
database.extra = json.dumps(
{
"engine_params": {
"connect_args": {"application": "Custom user agent", "foo": "bar"}
}
}
)
assert SnowflakeEngineSpec.get_extra_params(database) == {
"engine_params": {
"connect_args": {"application": "Custom user agent", "foo": "bar"}
}
}