ss-tools/backend/tests/plugins/test_migration_plugin_coverage.py

# #region Test.MigrationPlugin.Coverage [C:3] [TYPE Module] [SEMANTICS test,migration,coverage,edge]
# @BRIEF Coverage closure for Plugin.Migration.MigrationPlugin — archive helpers, composite-key
#       fallback failure branches, pre-migration rescan, password-path regex fallback, and
#       finally-block error paths. Hardcoded fixtures only; mocks stay at external boundaries.
# @RELATION BINDS_TO -> [Plugin.Migration.MigrationPlugin]
# @TEST_EDGE: empty_archive_inputs -> _try_save_failed_archive returns None without task_id/zip
# @TEST_EDGE: missing_live_target_contract -> source sync fails with no_live_target_contract
# @TEST_EDGE: re_transform_failure -> source fallback aborts when the retry transform fails
# @TEST_EDGE: invalid_mutation_server -> composite_key_mutation_server resets to "target"
# @TEST_EDGE: rescan_api_failure -> pre-migration ID rescan failure is non-fatal
# @TEST_EDGE: missing_db_env_rows -> replace_db_config tolerates absent Environment rows
# @TEST_EDGE: password_path_regex_fallback -> yaml path recovered from error text
# @TEST_EDGE: empty_passwords -> task without passwords fails gracefully, no pop
# @TEST_EDGE: export_failure_cleanup -> task password params cleaned on export failure
# @TEST_EDGE: lineage_hook_failure -> post-deploy lineage refresh failure is non-fatal
# @TEST_EDGE: engine_db_close_failure -> session close errors are swallowed
import pytest
from contextlib import ExitStack
from unittest.mock import AsyncMock, MagicMock, patch

from src.core.utils.network import SupersetAPIError
from src.models.storage import FileCategory
from src.plugins.migration import (
    MigrationPlugin,
    _failed_entry_with_archive,
    _try_save_failed_archive,
)


# ── Helpers ──


def _make_env(id_val="env-1", name="Source Env"):
    env = MagicMock()
    env.id = id_val
    env.name = name
    env.url = "https://superset.example.com"
    env.username = "admin"
    env.password = "secret"
    return env


def _make_dashboard(dash_id=1, title="Dash"):
    return {"id": dash_id, "slug": f"slug-{dash_id}", "dashboard_title": title}


def _make_superset_client():
    client = MagicMock()
    client.aclose = AsyncMock()
    return client


def _make_mapping_service():
    svc = MagicMock()
    svc.sync_environment = AsyncMock()
    return svc


def _base_mocks():
    """C1 helper: standard execute harness (environments, clients, engine)."""
    src_env = _make_env("env-1", "Source")
    tgt_env = _make_env("env-2", "Target")
    mock_cm = MagicMock()
    mock_cm.get_environments.return_value = [src_env, tgt_env]
    mock_src = _make_superset_client()
    mock_src.get_dashboards = AsyncMock(return_value=(True, [_make_dashboard(1, "Dash")]))
    mock_src.export_dashboard = AsyncMock(return_value=(b"zip", "meta"))
    mock_tgt = _make_superset_client()
    mock_tgt.import_dashboard = AsyncMock()
    mock_engine = MagicMock()
    mock_engine.transform_zip.return_value = True
    return mock_cm, mock_src, mock_tgt, mock_engine


def _patch_stack(
    mock_cm,
    mock_src,
    mock_tgt,
    mock_engine,
    *,
    tm=None,
    mapping_svc=None,
    session_locals=None,
    ctf_paths=None,
    extra_patches=None,
):
    """C1 helper: open the standard patch stack for MigrationPlugin.execute."""
    stack = ExitStack()
    stack.enter_context(patch("src.plugins.migration.get_config_manager", return_value=mock_cm))
    mock_sc = stack.enter_context(patch("src.plugins.migration.SupersetClient"))
    mock_sc.side_effect = [mock_src, mock_tgt]
    stack.enter_context(patch("src.plugins.migration.MigrationEngine", return_value=mock_engine))
    mock_ctf = stack.enter_context(patch("src.plugins.migration.create_temp_file"))
    if ctf_paths is None:
        mock_ctf.return_value.__enter__ = MagicMock(return_value="/tmp/test.zip")
    else:
        mock_ctf.return_value.__enter__ = MagicMock(side_effect=ctf_paths)
    if mapping_svc is None:
        mapping_svc = _make_mapping_service()
    stack.enter_context(patch("src.plugins.migration.IdMappingService", return_value=mapping_svc))
    session = stack.enter_context(patch("src.plugins.migration.SessionLocal"))
    if session_locals is not None:
        session.side_effect = session_locals
    if tm is not None:
        stack.enter_context(patch("src.dependencies.get_task_manager", return_value=tm))
    for target, new in (extra_patches or {}).items():
        stack.enter_context(patch(target, new=new))
    return stack


class TestMigrationArchiveHelpers:
    """Direct coverage for _try_save_failed_archive / _failed_entry_with_archive."""

    def test_try_save_skips_empty_inputs(self):
        assert _try_save_failed_archive(None, 1, "/tmp/x.zip") is None
        assert _try_save_failed_archive("task-1", 1, "") is None

    def test_try_save_persists_and_returns_rel(self):
        svc = MagicMock()
        with patch("src.plugins.migration.get_storage_service", return_value=svc):
            rel = _try_save_failed_archive("task-1", 7, "/tmp/export.zip")
        assert rel == "task-1/7.zip"
        svc.save_from_path.assert_called_once_with(
            FileCategory.MIGRATION_FAILED, "task-1/7.zip", "/tmp/export.zip"
        )

    def test_try_save_source_kind_suffix(self):
        svc = MagicMock()
        with patch("src.plugins.migration.get_storage_service", return_value=svc):
            rel = _try_save_failed_archive("task-1", 7, "/tmp/export.zip", kind="source")
        assert rel == "task-1/7.source.zip"

    def test_failed_entry_without_zip_paths(self):
        entry = _failed_entry_with_archive(
            1, "Dash", RuntimeError("boom"), phase="import", task_id="task-1"
        )
        assert entry["error"] == "boom"
        assert "source_archive_path" not in entry
        assert "archive_path" not in entry

    def test_failed_entry_persists_both_archives(self):
        svc = MagicMock()
        with patch("src.plugins.migration.get_storage_service", return_value=svc):
            entry = _failed_entry_with_archive(
                1, "Dash", RuntimeError("boom"), task_id="task-1",
                zip_path="/tmp/t.zip", source_zip_path="/tmp/s.zip",
            )
        assert entry["source_archive_path"] == "task-1/1.source.zip"
        assert entry["archive_path"] == "task-1/1.zip"


class TestMigrationCompositeKeyCoverage:
    """Composite-key fallback failure branches (source + target)."""

    @pytest.mark.asyncio
    async def test_source_fallback_no_live_target_contract(self):
        """Source sync: every source contract lacks a live target -> no_live_target_contract."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mock_engine.read_dataset_contracts_from_zip.return_value = [
            {"uuid": "ds-1", "database_uuid": "src-db", "catalog": None,
             "schema": "public", "table_name": "users"}
        ]
        mock_tgt.import_dashboard = AsyncMock(side_effect=RuntimeError("Generic import boom"))
        mock_read_live = AsyncMock(return_value={})
        mock_sync = AsyncMock(return_value={
            "changed": 0, "unchanged": 0, "skipped_missing": 0, "failed": 0, "errors": [],
        })

        with _patch_stack(
            mock_cm, mock_src, mock_tgt, mock_engine,
            extra_patches={
                "src.plugins.migration.read_live_dataset_contracts": mock_read_live,
                "src.plugins.migration.sync_dataset_composite_keys": mock_sync,
            },
        ):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "db_mappings": {"src-db": "tgt-db"},
                "sync_dataset_composite_keys": True,
                "composite_key_mutation_server": "source",
            })

        assert result["status"] == "PARTIAL_SUCCESS"
        entry = result["failed_dashboards"][0]
        assert entry["composite_key_retried"] is True
        assert entry["composite_key_sync_report"]["errors"] == [
            {"uuid": "ds-1", "error": "no_live_target_contract"}
        ]
        mock_read_live.assert_awaited_once()
        mock_sync.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_source_fallback_re_transform_failure(self, tmp_path):
        """Source sync succeeds but the re-transform fails -> fallback aborts."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        tmp_source = tmp_path / "source.zip"
        tmp_transformed = tmp_path / "transformed.zip"
        tmp_source.write_bytes(b"zip")
        mock_engine.read_dataset_contracts_from_zip.return_value = [
            {"uuid": "ds-1", "database_uuid": "src-db", "catalog": None,
             "schema": "public", "table_name": "users"}
        ]
        mock_engine.transform_zip.side_effect = [True, False]
        mock_src.export_dashboard = AsyncMock(side_effect=[(b"zip", "meta"), (b"zip2", "meta")])
        mock_tgt.import_dashboard = AsyncMock(side_effect=RuntimeError("Generic import boom"))
        mock_read_live = AsyncMock(return_value={
            "ds-1": {"uuid": "ds-1", "database_uuid": "tgt-db", "catalog": None,
                     "schema": "public", "table_name": "users"},
        })
        mock_sync = AsyncMock(return_value={
            "changed": 1, "unchanged": 0, "skipped_missing": 0, "failed": 0, "errors": [],
        })

        with _patch_stack(
            mock_cm, mock_src, mock_tgt, mock_engine,
            extra_patches={
                "src.plugins.migration.read_live_dataset_contracts": mock_read_live,
                "src.plugins.migration.sync_dataset_composite_keys": mock_sync,
            },
            ctf_paths=[str(tmp_source), str(tmp_transformed)],
        ):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "db_mappings": {"src-db": "tgt-db"},
                "sync_dataset_composite_keys": True,
                "composite_key_mutation_server": "source",
            })

        assert result["status"] == "PARTIAL_SUCCESS"
        entry = result["failed_dashboards"][0]
        assert entry["composite_key_retried"] is True
        assert entry["composite_key_sync_report"] == {
            "changed": 1, "unchanged": 0, "skipped_missing": 0, "failed": 0, "errors": [],
        }
        assert mock_engine.transform_zip.call_count == 2
        assert mock_src.export_dashboard.await_count == 2

    @pytest.mark.asyncio
    async def test_target_fallback_no_contracts_invalid_server(self):
        """Invalid mutation server resets to target; empty contracts abort target sync."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mock_engine.read_dataset_contracts_from_zip.return_value = []
        mock_tgt.import_dashboard = AsyncMock(side_effect=RuntimeError("Generic import boom"))
        mock_sync = AsyncMock(return_value={
            "changed": 0, "unchanged": 0, "skipped_missing": 0, "failed": 0, "errors": [],
        })

        with _patch_stack(
            mock_cm, mock_src, mock_tgt, mock_engine,
            extra_patches={"src.plugins.migration.sync_dataset_composite_keys": mock_sync},
        ):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "sync_dataset_composite_keys": True,
                "composite_key_mutation_server": "bogus",
            })

        assert result["status"] == "PARTIAL_SUCCESS"
        entry = result["failed_dashboards"][0]
        assert entry["composite_key_mutation_server"] == "target"
        assert entry["composite_key_retried"] is True
        mock_sync.assert_not_awaited()


class TestMigrationExecuteCoverage:
    """Remaining execute() branches: rescan, replace-block, password regex, cleanup."""

    @pytest.mark.asyncio
    async def test_rescan_success(self):
        """Pre-migration ID rescan succeeds and completes the full cycle."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mapping_svc = _make_mapping_service()

        with _patch_stack(mock_cm, mock_src, mock_tgt, mock_engine, mapping_svc=mapping_svc):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "fix_cross_filters": True, "rescan_ids_before_migration": True,
            })

        assert result["status"] == "SUCCESS"
        assert len(result["migrated_dashboards"]) == 1
        # rescan + post-migration incremental sync
        assert mapping_svc.sync_environment.await_count == 2

    @pytest.mark.asyncio
    async def test_rescan_failure_is_non_fatal(self):
        """Pre-migration rescan raising does not abort the batch."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mapping_svc = MagicMock()
        mapping_svc.sync_environment = AsyncMock(
            side_effect=[RuntimeError("rescan boom"), None]
        )

        with _patch_stack(mock_cm, mock_src, mock_tgt, mock_engine, mapping_svc=mapping_svc):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "fix_cross_filters": True, "rescan_ids_before_migration": True,
            })

        assert result["status"] == "SUCCESS"
        assert len(result["migrated_dashboards"]) == 1
        assert mapping_svc.sync_environment.await_count == 2

    @pytest.mark.asyncio
    async def test_replace_db_missing_env_rows_transform_fails_no_task(self):
        """Missing Environment rows skip stored mappings; transform failure without
        task_id never pauses for resolution."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mock_engine.transform_zip.return_value = False
        replace_db = MagicMock()
        replace_db.query.return_value.filter.return_value.first.return_value = None

        with _patch_stack(
            mock_cm, mock_src, mock_tgt, mock_engine,
            session_locals=[replace_db, MagicMock(), MagicMock(), MagicMock()],
        ):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": True,
            })

        assert result["status"] == "PARTIAL_SUCCESS"
        assert result["mapping_count"] == 0
        entry = result["failed_dashboards"][0]
        assert entry["phase"] == "transform"
        assert "Failed to transform ZIP" in entry["error"]
        assert mock_engine.transform_zip.call_count == 1

    @pytest.mark.asyncio
    async def test_password_path_regex_fallback(self):
        """Formatted error without password paths -> yaml path recovered from error text."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mock_tm = MagicMock()
        mock_tm.await_input = AsyncMock()
        mock_tm.wait_for_input = AsyncMock()
        task = MagicMock()
        task.params = {"passwords": {"SomeDB": "s3cret"}}
        mock_tm.get_task.return_value = task
        err = SupersetAPIError(
            "Superset import failed",
            status_code=422,
            errors=[{
                "message": "Must provide a password for the database to continue",
                "error_type": "DATABASE_NOT_FOUND",
            }],
            response_body='{"errors":[{"message":"Must provide a password for the database '
                          'to continue","extra":{"databases/SomeDB.yaml":{"message":"password '
                          'required"}}}]}',
        )
        mock_tgt.import_dashboard = AsyncMock(side_effect=[err, None])

        with _patch_stack(mock_cm, mock_src, mock_tgt, mock_engine, tm=mock_tm):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "_task_id": "task-pw-regex",
            })

        assert result["status"] == "SUCCESS"
        assert len(result["migrated_dashboards"]) == 1
        databases = mock_tm.await_input.await_args.args[1]["databases"]
        assert databases == ["SomeDB"]

    @pytest.mark.asyncio
    async def test_password_empty_passwords_skips_pop(self):
        """Task without a passwords payload fails gracefully and skips the pop."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mock_tm = MagicMock()
        mock_tm.await_input = AsyncMock()
        mock_tm.wait_for_input = AsyncMock()
        task = MagicMock()
        task.params = {}
        mock_tm.get_task.return_value = task
        mock_tgt.import_dashboard = AsyncMock(side_effect=RuntimeError(
            "Must provide a password for the database PostgreSQL"
        ))

        with _patch_stack(mock_cm, mock_src, mock_tgt, mock_engine, tm=mock_tm):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "_task_id": "task-pw-empty",
            })

        assert result["status"] == "PARTIAL_SUCCESS"
        entry = result["failed_dashboards"][0]
        assert entry["phase"] == "password_retry"

    @pytest.mark.asyncio
    async def test_export_failure_cleans_task_passwords(self):
        """Export-phase failure pops the passwords payload from the task."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mock_src.export_dashboard = AsyncMock(side_effect=RuntimeError("Export boom"))
        mock_tm = MagicMock()
        task = MagicMock()
        task.params = {"passwords": {"PG": "s3cret"}}
        mock_tm.get_task.return_value = task

        with _patch_stack(mock_cm, mock_src, mock_tgt, mock_engine, tm=mock_tm):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "_task_id": "task-exp-1",
            })

        assert result["status"] == "PARTIAL_SUCCESS"
        entry = result["failed_dashboards"][0]
        assert entry["phase"] == "export"
        assert "Export boom" in entry["error"]
        assert "passwords" not in task.params

    @pytest.mark.asyncio
    async def test_export_failure_task_without_passwords_key(self):
        """Export-phase failure with no passwords key is a no-op cleanup."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        mock_src.export_dashboard = AsyncMock(side_effect=RuntimeError("Export boom"))
        mock_tm = MagicMock()
        task = MagicMock()
        task.params = {}
        mock_tm.get_task.return_value = task

        with _patch_stack(mock_cm, mock_src, mock_tgt, mock_engine, tm=mock_tm):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
                "_task_id": "task-exp-2",
            })

        assert result["status"] == "PARTIAL_SUCCESS"
        assert result["failed_dashboards"][0]["phase"] == "export"

    @pytest.mark.asyncio
    async def test_lineage_hook_failure_non_fatal(self):
        """Post-deploy lineage refresh failure is logged, never fatal."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        lineage_sess = MagicMock()
        lineage_sess.query.side_effect = RuntimeError("lineage boom")

        with _patch_stack(
            mock_cm, mock_src, mock_tgt, mock_engine,
            session_locals=[MagicMock(), MagicMock(), lineage_sess],
        ):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
            })

        assert result["status"] == "SUCCESS"
        assert len(result["migrated_dashboards"]) == 1

    @pytest.mark.asyncio
    async def test_engine_db_close_failure_swallowed(self):
        """engine_db.close() raising is swallowed in the finally block."""
        mock_cm, mock_src, mock_tgt, mock_engine = _base_mocks()
        engine_db = MagicMock()
        engine_db.close.side_effect = RuntimeError("close boom")

        with _patch_stack(
            mock_cm, mock_src, mock_tgt, mock_engine,
            session_locals=[engine_db, MagicMock(), MagicMock()],
        ):
            result = await MigrationPlugin().execute({
                "source_env_id": "env-1", "target_env_id": "env-2",
                "selected_ids": [1], "replace_db_config": False,
            })

        assert result["status"] == "SUCCESS"
        assert len(result["migrated_dashboards"]) == 1
# #endregion Test.MigrationPlugin.Coverage