ss-tools/backend/tests/services/dataset_review/test_superset_matrix.py

# #region SupersetCompatibilityMatrixTests [C:2] [TYPE Module]
# @SEMANTICS: dataset_review, superset, compatibility_matrix, preview, sql_lab, tests
# @PURPOSE: Verifies Superset preview and SQL Lab endpoint fallback strategy used by dataset-review orchestration.
# @LAYER Tests
# @RELATION DEPENDS_ON ->[SupersetClient]
# @RELATION DEPENDS_ON ->[SupersetCompilationAdapter]

from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

import pytest

from src.core.utils.superset_compilation_adapter import (
    PreviewCompilationPayload,
    SqlLabLaunchPayload,
    SupersetCompilationAdapter,
)

# Import models to ensure proper SQLAlchemy registration


# #region make_adapter [C:2] [TYPE Function]
# @PURPOSE: Build an adapter with a mock Superset client and deterministic environment for compatibility tests.
# @RELATION DEPENDS_ON ->[SupersetCompilationAdapter]
def make_adapter():
    environment = SimpleNamespace(
        id="env-1",
        name="Test Env",
        url="http://superset.example",
        username="user",
        password="pass",
        verify_ssl=True,
        timeout=30,
    )
    client = MagicMock()
    client.network = MagicMock()
    return SupersetCompilationAdapter(environment=environment, client=client), client


# #endregion make_adapter


# #region test_preview_prefers_supported_client_method_before_network_fallback [C:2] [TYPE Function]
# @PURPOSE: Confirms preview compilation uses a supported client method first when the capability exists.
# @RELATION DEPENDS_ON ->[SupersetCompilationAdapter]
@pytest.mark.asyncio
async def test_preview_prefers_supported_client_method_before_network_fallback():
    adapter, client = make_adapter()
    client.compile_preview = AsyncMock(return_value={"compiled_sql": "SELECT 1"})
    payload = PreviewCompilationPayload(
        session_id="sess-1",
        dataset_id=42,
        preview_fingerprint="fp-1",
        template_params={"country": "RU"},
        effective_filters=[{"name": "country", "value": "RU"}],
    )

    preview = await adapter.compile_preview(payload)

    assert preview.preview_status.value == "ready"
    assert preview.compiled_sql == "SELECT 1"
    client.compile_preview.assert_called_once()
    client.network.request.assert_not_called()


# #endregion test_preview_prefers_supported_client_method_before_network_fallback


# #region test_preview_falls_back_across_matrix_until_supported_endpoint_returns_sql [C:2] [TYPE Function]
# @PURPOSE: Confirms preview fallback walks the compatibility matrix from preferred to legacy endpoints until one returns compiled SQL.
# @RELATION DEPENDS_ON ->[SupersetCompilationAdapter]
@pytest.mark.asyncio
async def test_preview_falls_back_across_matrix_until_supported_endpoint_returns_sql():
    adapter, client = make_adapter()
    payload = PreviewCompilationPayload(
        session_id="sess-2",
        dataset_id=77,
        preview_fingerprint="fp-2",
        template_params={"region": "emea"},
        effective_filters=[],
    )

    client.network.request = AsyncMock(side_effect=[
        RuntimeError("preview endpoint unavailable"),
        {"result": {"sql": "SELECT * FROM dataset_77"}},
    ])

    preview = await adapter.compile_preview(payload)

    assert preview.preview_status.value == "ready"
    assert preview.compiled_sql == "SELECT * FROM dataset_77"
    assert client.network.request.call_count == 2
    # @FRAGILE: Positional call assertion — ordering changes will break this test without indicating a real regression. Prefer content-based assertion.
    first_call = client.network.request.call_args_list[0].kwargs
    # @FRAGILE: Positional call assertion — ordering changes will break this test without indicating a real regression. Prefer content-based assertion.
    second_call = client.network.request.call_args_list[1].kwargs
    assert first_call["endpoint"] == "/dataset/77/preview"
    assert second_call["endpoint"] == "/dataset/77/sql"


# #endregion test_preview_falls_back_across_matrix_until_supported_endpoint_returns_sql


# #region test_sql_lab_launch_falls_back_to_legacy_execute_endpoint [C:2] [TYPE Function]
# @PURPOSE: Confirms SQL Lab launch falls back from modern to legacy execute endpoint and preserves canonical session reference extraction.
# @RELATION DEPENDS_ON ->[SupersetCompilationAdapter]
@pytest.mark.asyncio
async def test_sql_lab_launch_falls_back_to_legacy_execute_endpoint():
    adapter, client = make_adapter()
    client.get_dataset = AsyncMock(return_value={
        "result": {
            "id": 55,
            "schema": "public",
            "database": {"id": 9},
        }
    })
    client.network.request = AsyncMock(side_effect=[
        RuntimeError("sqllab execute unavailable"),
        {"result": {"id": "query-123"}},
    ])
    payload = SqlLabLaunchPayload(
        session_id="sess-3",
        dataset_id=55,
        preview_id="preview-9",
        compiled_sql="SELECT * FROM sales",
        template_params={"limit": 10},
    )

    sql_lab_ref = await adapter.create_sql_lab_session(payload)

    assert sql_lab_ref == "query-123"
    assert client.network.request.call_count == 2
    # @FRAGILE: Positional call assertion — ordering changes will break this test without indicating a real regression. Prefer content-based assertion.
    first_call = client.network.request.call_args_list[0].kwargs
    # @FRAGILE: Positional call assertion — ordering changes will break this test without indicating a real regression. Prefer content-based assertion.
    second_call = client.network.request.call_args_list[1].kwargs
    assert first_call["endpoint"] == "/sqllab/execute/"
    assert second_call["endpoint"] == "/sql_lab/execute/"


# #endregion test_sql_lab_launch_falls_back_to_legacy_execute_endpoint


# #endregion SupersetCompatibilityMatrixTests