ss-tools/backend/tests/test_datasets.py

# #region TestDatasetsApi [C:2] [TYPE Module]
# @RELATION  DEPENDS_ON -> [DatasetsApi]
# @SEMANTICS: tests, datasets, api, metrics, inline-edit
# @PURPOSE: Contract tests for dataset endpoints (list+stats, detail+metrics, inline-edit descriptions).
# @LAYER Tests
import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from fastapi.testclient import TestClient

from src.app import app
from src.dependencies import (
    get_config_manager,
    get_current_user,
    get_resource_service,
    get_task_manager,
    has_permission,
)

client = TestClient(app)

# #region mock_user [C:2] [TYPE Function]
# @PURPOSE: Create a mock user with Admin role and all permissions for testing.
def _make_mock_user():
    admin_role = MagicMock()
    admin_role.name = "Admin"
    admin_role.permissions = []

    mock_user = MagicMock()
    mock_user.username = "test-user"
    mock_user.roles = [admin_role]
    return mock_user


# #region mock_deps [C:2] [TYPE Function]
# @PURPOSE: Provide dependency override fixture for dataset route tests.
# @TEST_FIXTURE: dataset_route_overrides -> INLINE_JSON
@pytest.fixture
def mock_deps():
    config_manager = MagicMock()
    task_manager = MagicMock()
    resource_service = MagicMock()

    env = MagicMock()
    env.id = "ss-dev"
    env.name = "SS Dev"
    env.superset_url = "http://superset:8088"
    env.superset_token = "mock-token"
    config_manager.get_environments.return_value = [env]

    task_manager.get_all_tasks.return_value = []

    resource_service.get_datasets_with_status = AsyncMock(
        return_value=[
            {
                "id": 1,
                "table_name": "users",
                "schema": "public",
                "database": "postgres",
                "mapped_fields": {"total": 10, "mapped": 7},
                "last_task": {"task_id": "t1", "status": "SUCCESS"},
                "linked_dashboard_count": 2,
            },
            {
                "id": 2,
                "table_name": "orders",
                "schema": "public",
                "database": "postgres",
                "mapped_fields": {"total": 5, "mapped": 0},
                "last_task": None,
                "linked_dashboard_count": 0,
            },
            {
                "id": 3,
                "table_name": "products",
                "schema": "public",
                "database": "postgres",
                "mapped_fields": None,
                "last_task": None,
                "linked_dashboard_count": 0,
            },
        ]
    )

    mock_user = _make_mock_user()

    app.dependency_overrides[get_config_manager] = lambda: config_manager
    app.dependency_overrides[get_task_manager] = lambda: task_manager
    app.dependency_overrides[get_resource_service] = lambda: resource_service
    app.dependency_overrides[get_current_user] = lambda: mock_user
    # Override specific permission dependencies used by datasets routes
    app.dependency_overrides[has_permission("plugin:migration", "READ")] = lambda: mock_user
    app.dependency_overrides[has_permission("plugin:migration", "WRITE")] = lambda: mock_user

    yield

    app.dependency_overrides.clear()


# #region test_get_datasets_returns_stats [C:2] [TYPE Function]
# @PURPOSE: Verify GET /api/datasets returns stats object with correct counts.
# @TEST_CONTRACT: datasets_list -> datasets with stats payload
# @TEST_SCENARIO: datasets_with_stats -> HTTP 200 returns datasets, stats, and pagination.
# @TEST_INVARIANT: stats_counts_correct -> VERIFIED_BY: [datasets_with_stats]
def test_get_datasets_returns_stats(mock_deps):
    response = client.get(
        "/api/datasets?env_id=ss-dev&page=1&page_size=10"
    )
    assert response.status_code == 200
    data = response.json()
    assert "datasets" in data
    assert "stats" in data
    assert data["stats"]["total"] == 3
    assert data["stats"]["mapped_count"] == 1   # users: 7/10
    assert data["stats"]["unmapped_count"] == 1  # orders: 0/5
    assert data["stats"]["linked_count"] == 1    # users has 2 linked dashboards
    assert data["total"] == 3
    assert data["page"] == 1
    assert len(data["datasets"]) == 3


# #region test_get_datasets_filter_unmapped [C:2] [TYPE Function]
# @PURPOSE: Verify server-side filtering by unmapped datasets.
# @TEST_SCENARIO: filter_unmapped -> Only returns datasets with mapped=0.
def test_get_datasets_filter_unmapped(mock_deps):
    response = client.get(
        "/api/datasets?env_id=ss-dev&filter=unmapped&page=1&page_size=10"
    )
    assert response.status_code == 200
    data = response.json()
    # stats are global (unfiltered)
    assert data["stats"]["total"] == 3
    # filtered list: only orders (mapped=0)
    assert data["total"] == 1
    assert len(data["datasets"]) == 1
    assert data["datasets"][0]["table_name"] == "orders"


# #region test_get_datasets_filter_mapped [C:2] [TYPE Function]
# @PURPOSE: Verify server-side filtering by mapped datasets.
def test_get_datasets_filter_mapped(mock_deps):
    response = client.get(
        "/api/datasets?env_id=ss-dev&filter=mapped&page=1&page_size=10"
    )
    assert response.status_code == 200
    data = response.json()
    assert data["stats"]["total"] == 3
    assert data["total"] == 1
    assert data["datasets"][0]["table_name"] == "users"


# #region test_get_datasets_filter_linked [C:2] [TYPE Function]
# @PURPOSE: Verify server-side filtering by datasets linked to dashboards.
def test_get_datasets_filter_linked(mock_deps):
    response = client.get(
        "/api/datasets?env_id=ss-dev&filter=linked&page=1&page_size=10"
    )
    assert response.status_code == 200
    data = response.json()
    assert data["stats"]["total"] == 3
    assert data["total"] == 1
    assert data["datasets"][0]["table_name"] == "users"


# #region test_get_datasets_filter_all [C:2] [TYPE Function]
# @PURPOSE: Verify filter=all returns all datasets (same as no filter).
def test_get_datasets_filter_all(mock_deps):
    response = client.get(
        "/api/datasets?env_id=ss-dev&filter=all&page=1&page_size=10"
    )
    assert response.status_code == 200
    data = response.json()
    assert data["total"] == 3


# #region test_get_dataset_detail_returns_metrics [C:2] [TYPE Function]
# @PURPOSE: Verify GET /api/datasets/{id} returns metrics and metric_count.
# @TEST_SCENARIO: detail_with_metrics -> Response includes metrics list and metric_count.
def test_get_dataset_detail_returns_metrics(mock_deps):
    with patch("src.api.routes.datasets.SupersetClient") as MockClient:
        mock_client = MagicMock()
        MockClient.return_value = mock_client

        mock_client.get_dataset_detail.return_value = {
            "id": 1,
            "table_name": "users",
            "schema": "public",
            "database": "postgres",
            "description": "User table",
            "columns": [
                {"id": 10, "name": "id", "type": "INT", "is_dttm": False, "is_active": True, "description": "PK"},
                {"id": 11, "name": "name", "type": "STRING", "is_dttm": False, "is_active": True, "description": ""},
            ],
            "column_count": 2,
            "metrics": [
                {"id": 100, "metric_name": "count", "expression": "COUNT(*)", "verbose_name": "Row Count", "description": "Total rows", "metric_type": "count"},
            ],
            "metric_count": 1,
            "sql": "SELECT * FROM users",
            "linked_dashboards": [{"id": 5, "title": "Users Dashboard", "slug": "users-dash"}],
            "linked_dashboard_count": 1,
            "is_sqllab_view": False,
            "created_on": "2024-01-01T00:00:00",
            "changed_on": "2024-06-01T00:00:00",
        }

        response = client.get(
            "/api/datasets/1?env_id=ss-dev"
        )
        assert response.status_code == 200
        data = response.json()
        assert "metrics" in data
        assert "metric_count" in data
        assert data["metric_count"] == 1
        assert len(data["metrics"]) == 1
        assert data["metrics"][0]["metric_name"] == "count"
        assert data["metrics"][0]["expression"] == "COUNT(*)"


# #region test_update_column_description [C:2] [TYPE Function]
# @PURPOSE: Verify PUT /api/datasets/{id}/columns/{col_id}/description saves description.
# @TEST_SCENARIO: save_column_desc -> Saves and returns updated description.
def test_update_column_description(mock_deps):
    with patch("src.api.routes.datasets.SupersetClient") as MockClient:
        mock_client = MagicMock()
        MockClient.return_value = mock_client

        mock_client.get_dataset.return_value = {
            "id": 1,
            "result": {
                "id": 1,
                "columns": [
                    {"id": 10, "column_name": "id", "type": "INT", "description": "old desc"},
                    {"id": 11, "column_name": "name", "type": "STRING", "description": ""},
                ],
                "metrics": [],
            },
        }
        mock_client.update_dataset.return_value = {"id": 1}

        response = client.put(
            "/api/datasets/1/columns/10/description?env_id=ss-dev",
            json={"description": "new description"},
        )
        assert response.status_code == 200
        data = response.json()
        assert data["success"] is True
        assert data["description"] == "new description"

        # Verify update_dataset was called with correct column modified
        update_call = mock_client.update_dataset.call_args
        assert update_call[0][0] == 1  # dataset_id
        # Check override_columns=false was passed
        assert update_call[1]["override_columns"] is False


# #region test_update_metric_description [C:2] [TYPE Function]
# @PURPOSE: Verify PUT /api/datasets/{id}/metrics/{metric_id}/description saves correctly.
# @TEST_SCENARIO: save_metric_desc -> Saves and returns updated description.
def test_update_metric_description(mock_deps):
    with patch("src.api.routes.datasets.SupersetClient") as MockClient:
        mock_client = MagicMock()
        MockClient.return_value = mock_client

        mock_client.get_dataset.return_value = {
            "id": 1,
            "result": {
                "id": 1,
                "columns": [],
                "metrics": [
                    {"id": 100, "metric_name": "count", "expression": "COUNT(*)", "description": ""},
                ],
            },
        }
        mock_client.update_dataset.return_value = {"id": 1}

        response = client.put(
            "/api/datasets/1/metrics/100/description?env_id=ss-dev",
            json={"description": "Total row count"},
        )
        assert response.status_code == 200
        data = response.json()
        assert data["success"] is True
        assert data["description"] == "Total row count"


# #region test_update_column_description_not_found [C:2] [TYPE Function]
# @PURPOSE: Verify 404 when column not found.
def test_update_column_description_not_found(mock_deps):
    with patch("src.api.routes.datasets.SupersetClient") as MockClient:
        mock_client = MagicMock()
        MockClient.return_value = mock_client

        mock_client.get_dataset.return_value = {
            "id": 1,
            "result": {
                "id": 1,
                "columns": [{"id": 10, "column_name": "id"}],
                "metrics": [],
            },
        }

        response = client.put(
            "/api/datasets/1/columns/999/description?env_id=ss-dev",
            json={"description": "test"},
        )
        assert response.status_code == 404


# #region test_update_column_description_validation [C:2] [TYPE Function]
# @PURPOSE: Verify 400 when description exceeds max length.
def test_update_column_description_validation(mock_deps):
    response = client.put(
        "/api/datasets/1/columns/10/description?env_id=ss-dev",
        json={"description": "x" * 2001},
    )
    assert response.status_code == 400


# #region test_get_datasets_superset_503 [C:2] [TYPE Function]
# @PURPOSE: Verify GET /api/datasets returns 503 when Superset/ResourceService throws.
# @TEST_SCENARIO: upstream_superset_failure -> 503 with error detail.
# @TEST_INVARIANT: upstream_failure_503 -> VERIFIED_BY: [superset_failure]
def test_get_datasets_superset_503(mock_deps):
    # Override resource_service to throw during get_datasets_with_status
    from src.dependencies import get_resource_service

    failing_resource_service = MagicMock()
    failing_resource_service.get_datasets_with_status = AsyncMock(
        side_effect=Exception("Superset connection refused")
    )

    app.dependency_overrides[get_resource_service] = lambda: failing_resource_service

    try:
        response = client.get(
            "/api/datasets?env_id=ss-dev&page=1&page_size=10"
        )
        assert response.status_code == 503
        assert "Superset" in response.text or "Failed" in response.text
    finally:
        app.dependency_overrides[get_resource_service] = None


# #region test_update_column_description_superset_502 [C:2] [TYPE Function]
# @PURPOSE: Verify PUT column description returns 502 when SupersetClient.update_dataset throws.
# @TEST_SCENARIO: superset_put_failure -> 502 with error detail.
def test_update_column_description_superset_502(mock_deps):
    with patch("src.api.routes.datasets.SupersetClient") as MockClient:
        mock_client = MagicMock()
        MockClient.return_value = mock_client

        # GET succeeds
        mock_client.get_dataset.return_value = {
            "id": 1,
            "result": {
                "id": 1,
                "columns": [{"id": 10, "column_name": "id", "type": "INT", "description": "old desc"}],
                "metrics": [],
            },
        }
        # PUT fails
        mock_client.update_dataset.side_effect = Exception("Superset write failure")

        response = client.put(
            "/api/datasets/1/columns/10/description?env_id=ss-dev",
            json={"description": "new description"},
        )
        assert response.status_code == 502
        data = response.json()
        assert "Superset" in data.get("detail", "")


# #region test_update_column_description_strips_html [C:2] [TYPE Function]
# @PURPOSE: Verify HTML tags are stripped from description before saving.
# @TEST_SCENARIO: html_in_description -> HTML tags stripped, clean text saved.
def test_update_column_description_strips_html(mock_deps):
    with patch("src.api.routes.datasets.SupersetClient") as MockClient:
        mock_client = MagicMock()
        MockClient.return_value = mock_client

        mock_client.get_dataset.return_value = {
            "id": 1,
            "result": {
                "id": 1,
                "columns": [{"id": 10, "column_name": "id", "type": "INT", "description": ""}],
                "metrics": [],
            },
        }
        mock_client.update_dataset.return_value = {"id": 1}

        response = client.put(
            "/api/datasets/1/columns/10/description?env_id=ss-dev",
            json={"description": "Hello <b>World</b> with <script>alert('xss')</script>"},
        )
        assert response.status_code == 200
        data = response.json()
        # HTML stripped
        assert data["description"] == "Hello World with alert('xss')"
        assert "<b>" not in data["description"]
        assert "<script>" not in data["description"]

        # Verify the stripped description was sent to Superset
        update_call = mock_client.update_dataset.call_args
        updated_columns = update_call[0][1]["columns"]
        assert updated_columns[0]["description"] == "Hello World with alert('xss')"


# #region test_dataset_item_has_metric_count [C:2] [TYPE Function]
# @PURPOSE: Verify DatasetItem in list response includes metric_count field.
def test_dataset_item_has_metric_count(mock_deps):
    response = client.get(
        "/api/datasets?env_id=ss-dev&page=1&page_size=10"
    )
    assert response.status_code == 200
    data = response.json()
    for ds in data["datasets"]:
        assert "metric_count" in ds
        assert isinstance(ds["metric_count"], int)