558 lines
22 KiB
Python
558 lines
22 KiB
Python
# #region Test.SupersetSqlLabExecutor [C:3] [TYPE Module] [SEMANTICS test, superset, sqllab, executor, poll]
|
|
# @BRIEF Verify SupersetSqlLabExecutor contracts — _get_client, resolve_database_id, execute_sql,
|
|
# poll_execution_status, execute_and_poll, get_query_results.
|
|
# @RELATION BINDS_TO -> [SupersetSqlLabExecutor]
|
|
# @TEST_EDGE: env_not_found -> Raises ValueError
|
|
# @TEST_EDGE: target_database_id_resolution -> Resolves via target_database_id
|
|
# @TEST_EDGE: database_name_resolution -> Resolves by database_name
|
|
# @TEST_EDGE: databases_empty -> Raises ValueError
|
|
# @TEST_EDGE: database_name_not_found -> Raises ValueError
|
|
# @TEST_EDGE: execute_sql_no_database_id -> Falls back to resolve_database_id
|
|
# @TEST_EDGE: execute_sql_missing_query_id -> Returns result with None query_id
|
|
# @TEST_EDGE: polling_success -> Returns success status
|
|
# @TEST_EDGE: polling_failed -> Returns failed status
|
|
# @TEST_EDGE: polling_timeout -> Returns timeout status
|
|
# @TEST_EDGE: polling_error_retry -> Retries on network error
|
|
# @TEST_EDGE: execute_and_poll_sync_success -> Returns directly from sync mode
|
|
# @TEST_EDGE: execute_and_poll_no_query_id -> Returns failed
|
|
# @TEST_EDGE: get_query_results_success -> Returns results
|
|
# @TEST_EDGE: get_query_results_failure -> Returns error
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from src.core.config_manager import ConfigManager
|
|
from src.plugins.translate.superset_executor import SupersetSqlLabExecutor
|
|
|
|
|
|
def _make_config_manager(environments=None):
|
|
"""Create a mock ConfigManager with the given environments."""
|
|
config = MagicMock(spec=ConfigManager)
|
|
if environments is not None:
|
|
config.get_environments.return_value = environments
|
|
for env in environments:
|
|
config.get_environment.return_value = env
|
|
else:
|
|
config.get_environments.return_value = []
|
|
config.get_environment.return_value = None
|
|
return config
|
|
|
|
|
|
def _make_env(id="env-1", name="Test Env"):
|
|
env = MagicMock()
|
|
env.id = id
|
|
env.name = name
|
|
return env
|
|
|
|
|
|
class TestInit:
|
|
"""Verify initialization."""
|
|
|
|
def test_init(self):
|
|
config = _make_config_manager()
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
assert executor.config_manager is config
|
|
assert executor.env_id == "env-1"
|
|
assert executor._client is None
|
|
assert executor._database_id is None
|
|
|
|
|
|
class TestGetClient:
|
|
"""Verify _get_client."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lazy_initialization(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
with patch('src.plugins.translate.superset_executor.get_superset_client',
|
|
new=AsyncMock(return_value=MagicMock())) as mock_get:
|
|
client = await executor._get_client()
|
|
assert client is not None
|
|
assert executor._client is not None
|
|
# Second call returns cached client
|
|
client2 = await executor._get_client()
|
|
assert client2 is client
|
|
assert mock_get.call_count == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_env_not_found(self):
|
|
config = _make_config_manager([]) # no environments
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
with pytest.raises(ValueError, match="not found"):
|
|
await executor._get_client()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_env_not_found_empty_list(self):
|
|
config = _make_config_manager()
|
|
config.get_environments.return_value = []
|
|
executor = SupersetSqlLabExecutor(config, "nonexistent")
|
|
with pytest.raises(ValueError, match="not found"):
|
|
await executor._get_client()
|
|
|
|
|
|
class TestResolveDatabaseId:
|
|
"""Verify resolve_database_id."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_by_target_database_id(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
mock_client = AsyncMock()
|
|
mock_client.get_database.return_value = {
|
|
"result": {"id": 42, "backend": "postgresql", "database_name": "Analytics DB"},
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
db_id = await executor.resolve_database_id(target_database_id="42")
|
|
assert db_id == 42
|
|
assert executor._database_backend == "postgresql"
|
|
assert executor._database_name == "Analytics DB"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_by_target_database_id_fallback(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
mock_client = AsyncMock()
|
|
# get_database raises error, fallback to database list
|
|
mock_client.get_database.side_effect = ValueError("API error")
|
|
mock_client.get_databases.return_value = (None, [
|
|
{"id": 1, "database_name": "Main DB", "backend": "postgresql"},
|
|
])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
db_id = await executor.resolve_database_id(target_database_id="99")
|
|
assert db_id == 1
|
|
assert executor._database_backend == "postgresql"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_by_database_name(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
mock_client = AsyncMock()
|
|
mock_client.get_databases.return_value = (None, [
|
|
{"id": 1, "database_name": "Main DB", "backend": "postgresql"},
|
|
{"id": 2, "database_name": "Analytics DB", "backend": "clickhouse"},
|
|
])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
db_id = await executor.resolve_database_id(database_name="Analytics DB")
|
|
assert db_id == 2
|
|
assert executor._database_backend == "clickhouse"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_by_database_name_case_insensitive(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
mock_client = AsyncMock()
|
|
mock_client.get_databases.return_value = (None, [
|
|
{"id": 1, "database_name": "Main DB", "backend": "postgresql"},
|
|
])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
db_id = await executor.resolve_database_id(database_name="main db")
|
|
assert db_id == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_database_name_not_found(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
mock_client = AsyncMock()
|
|
mock_client.get_databases.return_value = (None, [
|
|
{"id": 1, "database_name": "Main DB", "backend": "postgresql"},
|
|
])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
with pytest.raises(ValueError, match="not found"):
|
|
await executor.resolve_database_id(database_name="Nonexistent DB")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_default_first_database(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
mock_client = AsyncMock()
|
|
mock_client.get_databases.return_value = (None, [
|
|
{"id": 42, "database_name": "Default DB", "backend": "mysql"},
|
|
])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
db_id = await executor.resolve_database_id()
|
|
assert db_id == 42
|
|
assert executor._database_backend == "mysql"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_no_databases_raises(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
mock_client = AsyncMock()
|
|
mock_client.get_databases.return_value = (None, [])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
with pytest.raises(ValueError, match="No databases found"):
|
|
await executor.resolve_database_id()
|
|
|
|
|
|
class TestGetDatabaseBackend:
|
|
"""Verify get_database_backend."""
|
|
|
|
def test_cached_value(self):
|
|
config = _make_config_manager()
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
assert executor.get_database_backend() is None
|
|
executor._database_backend = "postgresql"
|
|
assert executor.get_database_backend() == "postgresql"
|
|
|
|
|
|
class TestGetDatabaseName:
|
|
"""Verify get_database_name."""
|
|
|
|
def test_cached_value(self):
|
|
config = _make_config_manager()
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
assert executor.get_database_name() is None
|
|
executor._database_name = "Main DB"
|
|
assert executor.get_database_name() == "Main DB"
|
|
|
|
|
|
class TestExecuteSql:
|
|
"""Verify execute_sql."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_with_database_id(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
executor._database_id = 42
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"query_id": "q-123",
|
|
"status": "success",
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.execute_sql("SELECT 1", database_id=42)
|
|
assert result["query_id"] == "q-123"
|
|
assert result["status"] == "success"
|
|
assert result["database_id"] == 42
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_falls_back_to_resolve(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"query_id": "q-456",
|
|
"status": "success",
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
with patch.object(executor, 'resolve_database_id',
|
|
new=AsyncMock(return_value=99)):
|
|
result = await executor.execute_sql("SELECT 1")
|
|
assert result["query_id"] == "q-456"
|
|
assert result["database_id"] == 99
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_api_error(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
executor._database_id = 42
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.side_effect = ValueError("Connection refused")
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
with pytest.raises(ValueError, match="Superset SQL Lab execute failed"):
|
|
await executor.execute_sql("SELECT 1")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_missing_query_id(self):
|
|
"""Response missing query_id — still returns result with None query_id."""
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
executor._database_id = 42
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {"status": "pending"}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.execute_sql("SELECT 1")
|
|
assert result["query_id"] is None
|
|
assert result["status"] == "pending"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_query_id_in_result(self):
|
|
"""Query ID nested in result block."""
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
executor._database_id = 42
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"result": {"query_id": "q-nested"},
|
|
"status": "success",
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.execute_sql("SELECT 1")
|
|
assert result["query_id"] == "q-nested"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_query_id_in_query_block(self):
|
|
"""Query ID nested in query block."""
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
executor._database_id = 42
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"query": {"id": "q-from-query"},
|
|
"status": "pending",
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.execute_sql("SELECT 1")
|
|
assert result["query_id"] == "q-from-query"
|
|
|
|
|
|
class TestPollExecutionStatus:
|
|
"""Verify poll_execution_status."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_success(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"result": {"status": "success", "state": "success", "rows": 10},
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.poll_execution_status("q-1", max_polls=1)
|
|
assert result["status"] == "success"
|
|
assert result["rows_affected"] == 10
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_failed(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"result": {"status": "failed", "state": "failed", "error_message": "Query error"},
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.poll_execution_status("q-1", max_polls=1)
|
|
assert result["status"] == "failed"
|
|
assert "Query error" in result["error_message"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_timeout(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
# Always return pending
|
|
mock_client.network.request.return_value = {
|
|
"result": {"status": "pending", "state": "pending"},
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.poll_execution_status("q-1", max_polls=1, poll_interval_seconds=0.01)
|
|
assert result["status"] == "timeout"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_network_error_retry(self):
|
|
"""Network error should be retried."""
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
# First call fails, second succeeds
|
|
mock_client.network.request = AsyncMock(side_effect=[
|
|
ValueError("Network error"),
|
|
{"result": {"status": "success", "state": "success", "rows": 5}},
|
|
])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.poll_execution_status("q-1", max_polls=5, poll_interval_seconds=0.01)
|
|
assert result["status"] == "success"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_unknown_state_continues(self):
|
|
"""Unknown state should continue polling."""
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request = AsyncMock(side_effect=[
|
|
{"result": {"status": "weird_state", "state": "weird_state"}},
|
|
{"result": {"status": "success", "state": "success", "rows": 3}},
|
|
])
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.poll_execution_status("q-1", max_polls=5, poll_interval_seconds=0.01)
|
|
assert result["status"] == "success"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_uses_raw_result_when_no_result_key(self):
|
|
"""When response has no 'result' key, use raw dict."""
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"status": "success",
|
|
"state": "success",
|
|
"rows": 7,
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.poll_execution_status("q-1", max_polls=1)
|
|
assert result["status"] == "success"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_not_dict_response(self):
|
|
"""When response is not a dict, use empty dict."""
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = "not a dict"
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.poll_execution_status("q-1", max_polls=5, poll_interval_seconds=0.01)
|
|
assert result["status"] == "timeout"
|
|
|
|
|
|
class TestExecuteAndPoll:
|
|
"""Verify execute_and_poll."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_mode_success(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
with patch.object(executor, 'execute_sql',
|
|
new=AsyncMock(return_value={
|
|
"query_id": "q-1",
|
|
"status": "success",
|
|
"raw_response": {
|
|
"query": {"rows": 5, "endDttm": "2024-01-01"},
|
|
"data": [{"col": "val"}],
|
|
},
|
|
})):
|
|
result = await executor.execute_and_poll("SELECT 1")
|
|
assert result["status"] == "success"
|
|
assert result["rows_affected"] == 5
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_mode_poll(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
with patch.object(executor, 'execute_sql',
|
|
new=AsyncMock(return_value={
|
|
"query_id": "q-1",
|
|
"status": "pending",
|
|
"raw_response": {},
|
|
})):
|
|
with patch.object(executor, 'poll_execution_status',
|
|
new=AsyncMock(return_value={
|
|
"query_id": "q-1",
|
|
"status": "success",
|
|
"state": "success",
|
|
})):
|
|
result = await executor.execute_and_poll("SELECT 1")
|
|
assert result["status"] == "success"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_query_id_returns_failed(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
with patch.object(executor, 'execute_sql',
|
|
new=AsyncMock(return_value={
|
|
"query_id": None,
|
|
"status": "pending",
|
|
"raw_response": {},
|
|
})):
|
|
result = await executor.execute_and_poll("SELECT 1")
|
|
assert result["status"] == "failed"
|
|
assert "No query_id" in result["error_message"]
|
|
|
|
|
|
class TestGetQueryResults:
|
|
"""Verify get_query_results."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_success(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {
|
|
"results": [{"col": "val"}],
|
|
}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.get_query_results("q-1")
|
|
assert result["status"] == "success"
|
|
assert result["results"] == [{"col": "val"}]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_results(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.return_value = {}
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.get_query_results("q-1")
|
|
assert result["status"] == "no_results"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_error(self):
|
|
env = _make_env()
|
|
config = _make_config_manager([env])
|
|
executor = SupersetSqlLabExecutor(config, "env-1")
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.network.request.side_effect = ValueError("Not found")
|
|
|
|
with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
|
|
result = await executor.get_query_results("q-1")
|
|
assert result["status"] == "failed"
|
|
assert "Not found" in result["error_message"]
|
|
# #endregion Test.SupersetSqlLabExecutor
|