feat(db_engine): Add impersonation tests for StarRocks

Author: Woellchen

tests/unit_tests/db_engine_specs/test_starrocks.py

@@ -18,6 +18,7 @@
 from typing import Any, Optional
 
 import pytest
+from pytest_mock import MockerFixture
 from sqlalchemy import JSON, types
 from sqlalchemy.engine.url import make_url
 
@@ -124,3 +125,47 @@ def test_get_schema_from_engine_params() -> None:
         )
         is None
     )
+
+
+def test_impersonation_username(mocker: MockerFixture) -> None:
+    """
+    Test impersonation and make sure that `get_url_for_impersonation` leaves the URL
+    unchanged and that `get_prequeries` returns the appropriate impersonation query.
+    """
+    from superset.db_engine_specs.starrocks import StarRocksEngineSpec
+
+    database = mocker.MagicMock()
+    database.impersonate_user = True
+    database.get_effective_user.return_value = "alice"
+
+    assert StarRocksEngineSpec.get_url_for_impersonation(
+        url=make_url("starrocks://service_user@localhost:9030/hive.default"),
+        impersonate_user=True,
+        username="alice",
+        access_token=None,
+    ) == make_url("starrocks://service_user@localhost:9030/hive.default")
+
+    assert StarRocksEngineSpec.get_prequeries(database) == [
+        'EXECUTE AS "alice" WITH NO REVERT;'
+    ]
+
+
+def test_impersonation_disabled(mocker: MockerFixture) -> None:
+    """
+    Test that impersonation is not applied when the feature is disabled in
+    `get_url_for_impersonation` and `get_prequeries`.
+    """
+    from superset.db_engine_specs.starrocks import StarRocksEngineSpec
+
+    database = mocker.MagicMock()
+    database.impersonate_user = False
+    database.get_effective_user.return_value = "alice"
+
+    assert StarRocksEngineSpec.get_url_for_impersonation(
+        url=make_url("starrocks://service_user@localhost:9030/hive.default"),
+        impersonate_user=False,
+        username="alice",
+        access_token=None,
+    ) == make_url("starrocks://service_user@localhost:9030/hive.default")
+
+    assert StarRocksEngineSpec.get_prequeries(database) == []