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

# #region Test.DatasetReview.Orchestrator [C:3] [TYPE Module] [SEMANTICS test,dataset,orchestrator,session,launch,preview]
# @BRIEF Tests for orchestrator.py — DatasetReviewOrchestrator: start_session, prepare_launch_preview,
#        launch_dataset, _build_recovery_bootstrap, _enqueue_recovery_task.
# @RELATION BINDS_TO -> [DatasetReviewOrchestrator]
# @TEST_EDGE: start_session -> Superset link, dataset selection, validation errors
# @TEST_EDGE: prepare_launch_preview -> Success, missing session, preview blockers
# @TEST_EDGE: launch_dataset -> Success, failed SQL Lab, launch blockers
# @TEST_EDGE: _build_recovery_bootstrap -> Filter recovery, template variable discovery, execution mappings
# @TEST_EDGE: _enqueue_recovery_task -> Task linked, task manager absent, incompatible contract

from __future__ import annotations
from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))

from datetime import datetime, timezone
from unittest.mock import MagicMock, AsyncMock, patch, PropertyMock, call
import pytest


# ── Fixtures ────────────────────────────────────────────────

@pytest.fixture
def mock_repository():
    """Mock DatasetReviewSessionRepository."""
    repo = MagicMock()
    repo.db = MagicMock()
    repo.event_logger = MagicMock()
    repo.bump_session_version = MagicMock(return_value=2)
    return repo


@pytest.fixture
def mock_config_manager():
    """Mock ConfigManager.get_environment."""
    cm = MagicMock()
    env = MagicMock()
    env.id = "env-1"
    env.name = "Test Environment"
    cm.get_environment.return_value = env
    return cm


@pytest.fixture
def mock_user():
    from src.models.auth import User
    user = MagicMock(spec=User)
    user.id = "user-1"
    return user


@pytest.fixture
def orchestrator(mock_repository, mock_config_manager):
    from src.services.dataset_review.orchestrator import DatasetReviewOrchestrator
    return DatasetReviewOrchestrator(
        repository=mock_repository,
        config_manager=mock_config_manager,
    )


# ── Tests for start_session ────────────────────────────────

class TestStartSession:
    """start_session — session creation with validation."""

    @pytest.mark.asyncio
    async def test_start_session_dataset_selection(self, orchestrator, mock_repository, mock_user):
        """Dataset selection creates session with REVIEW_READY state."""
        from src.services.dataset_review.orchestrator_pkg._commands import StartSessionCommand

        command = StartSessionCommand(
            source_kind="dataset_selection",
            source_input="42",
            environment_id="env-1",
            user=mock_user,
        )

        # Setup repository mock chain
        persisted_session = MagicMock()
        persisted_session.session_id = "sess-1"
        persisted_session.current_phase = MagicMock(value="review")
        persisted_session.readiness_state = MagicMock(value="review_ready")
        persisted_session.source_kind = "dataset_selection"
        persisted_session.dataset_ref = "dataset:42"
        persisted_session.dataset_id = 42
        persisted_session.dashboard_id = None
        persisted_session.active_task_id = None
        persisted_session.user_id = "user-1"
        mock_repository.create_session.return_value = persisted_session
        mock_repository.save_profile_and_findings.return_value = persisted_session
        mock_repository.save_recovery_state.return_value = persisted_session

        with patch("src.services.dataset_review.orchestrator.build_initial_profile") as mock_build_profile, \
             patch("src.services.dataset_review.orchestrator.parse_dataset_selection", return_value=("dataset:42", 42)):
            mock_profile = MagicMock()
            mock_build_profile.return_value = mock_profile

            result = await orchestrator.start_session(command)

        assert result.session is not None
        assert result.session.session_id == "sess-1"
        mock_repository.create_session.assert_called_once()
        mock_repository.save_profile_and_findings.assert_called_once()
        mock_repository.event_logger.log_event.assert_called()

    @pytest.mark.asyncio
    async def test_start_session_empty_source_raises(self, orchestrator, mock_user):
        """Empty source_input raises ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import StartSessionCommand

        command = StartSessionCommand(
            source_kind="dataset_selection",
            source_input="",
            environment_id="env-1",
            user=mock_user,
        )

        with pytest.raises(ValueError, match="non-empty"):
            await orchestrator.start_session(command)

    @pytest.mark.asyncio
    async def test_start_session_invalid_source_kind_raises(self, orchestrator, mock_user):
        """Invalid source_kind raises ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import StartSessionCommand

        command = StartSessionCommand(
            source_kind="invalid_kind",
            source_input="test",
            environment_id="env-1",
            user=mock_user,
        )

        with pytest.raises(ValueError, match="source_kind must be"):
            await orchestrator.start_session(command)

    @pytest.mark.asyncio
    async def test_start_session_environment_not_found(self, orchestrator, mock_config_manager, mock_user):
        """Missing environment raises ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import StartSessionCommand
        mock_config_manager.get_environment.return_value = None

        command = StartSessionCommand(
            source_kind="dataset_selection",
            source_input="42",
            environment_id="nonexistent",
            user=mock_user,
        )

        with pytest.raises(ValueError, match="Environment not found"):
            await orchestrator.start_session(command)

    @pytest.mark.asyncio
    async def test_start_session_superset_link_partial_recovery(self, orchestrator, mock_repository, mock_config_manager, mock_user):
        """Superset link with partial recovery creates RECOVERY_REQUIRED session."""
        from src.services.dataset_review.orchestrator_pkg._commands import StartSessionCommand

        # Mock parsed context
        parsed_context = MagicMock()
        parsed_context.partial_recovery = True
        parsed_context.unresolved_references = ["dataset_column_missing"]
        parsed_context.dataset_ref = "schema.table"
        parsed_context.dataset_id = 42
        parsed_context.dashboard_id = 99
        parsed_context.dataset_payload = {"id": 42}

        command = StartSessionCommand(
            source_kind="superset_link",
            source_input="http://superset.example.com/dashboard/99",
            environment_id="env-1",
            user=mock_user,
        )

        persisted_session = MagicMock()
        persisted_session.session_id = "sess-2"
        persisted_session.current_phase = MagicMock(value="recovery")
        persisted_session.readiness_state = MagicMock(value="recovery_required")
        persisted_session.source_kind = "superset_link"
        persisted_session.dataset_ref = "schema.table"
        persisted_session.dataset_id = 42
        persisted_session.dashboard_id = 99
        persisted_session.active_task_id = None
        persisted_session.user_id = "user-1"
        mock_repository.create_session.return_value = persisted_session
        mock_repository.save_profile_and_findings.return_value = persisted_session
        mock_repository.save_recovery_state.return_value = persisted_session

        with patch("src.services.dataset_review.orchestrator.SupersetContextExtractor") as mock_extractor_cls, \
             patch("src.services.dataset_review.orchestrator.build_initial_profile") as mock_build_profile, \
             patch("src.services.dataset_review.orchestrator.build_partial_recovery_findings") as mock_build_findings:

            mock_extractor = MagicMock()
            mock_extractor.parse_superset_link = AsyncMock(return_value=parsed_context)
            mock_extractor.recover_imported_filters = AsyncMock(return_value=[])
            mock_extractor.discover_template_variables.return_value = []
            mock_extractor_cls.return_value = mock_extractor

            mock_profile = MagicMock()
            mock_build_profile.return_value = mock_profile
            mock_build_findings.return_value = [MagicMock()]

            result = await orchestrator.start_session(command)

        assert result is not None
        assert result.session.session_id == "sess-2"

    @pytest.mark.asyncio
    async def test_start_session_with_task_manager(self, orchestrator, mock_repository, mock_config_manager, mock_user):
        """Session start with task manager enqueues recovery task."""
        from src.services.dataset_review.orchestrator_pkg._commands import StartSessionCommand

        # Attach task manager
        mock_task_manager = MagicMock()
        task_object = MagicMock()
        task_object.id = "task-1"
        mock_task_manager.create_task.return_value = task_object
        orchestrator.task_manager = mock_task_manager

        command = StartSessionCommand(
            source_kind="dataset_selection",
            source_input="42",
            environment_id="env-1",
            user=mock_user,
        )

        persisted_session = MagicMock()
        persisted_session.session_id = "sess-3"
        persisted_session.current_phase = MagicMock(value="review")
        persisted_session.readiness_state = MagicMock(value="review_ready")
        persisted_session.source_kind = "dataset_selection"
        persisted_session.dataset_ref = "dataset:42"
        persisted_session.dataset_id = 42
        persisted_session.dashboard_id = None
        persisted_session.active_task_id = None
        persisted_session.user_id = "user-1"
        persisted_session.environment_id = "env-1"
        persisted_session.source_input = "42"
        mock_repository.create_session.return_value = persisted_session
        mock_repository.save_profile_and_findings.return_value = persisted_session

        with patch("src.services.dataset_review.orchestrator.build_initial_profile") as mock_build_profile, \
             patch("src.services.dataset_review.orchestrator.parse_dataset_selection", return_value=("dataset:42", 42)):
            mock_profile = MagicMock()
            mock_build_profile.return_value = mock_profile

            result = await orchestrator.start_session(command)

        assert result.session.active_task_id == "task-1"
        mock_repository.bump_session_version.assert_called()
        mock_repository.db.commit.assert_called()
        mock_repository.db.refresh.assert_called()


# ── Tests for prepare_launch_preview ───────────────────────

class TestPrepareLaunchPreview:
    """prepare_launch_preview — preview compilation."""

    def test_prepare_launch_preview_success(self, orchestrator, mock_repository, mock_config_manager):
        """Successful preview compilation returns PreparePreviewResult."""
        from src.services.dataset_review.orchestrator_pkg._commands import PreparePreviewCommand
        from src.services.dataset_review.orchestrator_pkg._helpers import PreviewStatus

        session = MagicMock()
        session.session_id = "sess-1"
        session.user_id = "user-1"
        session.dataset_id = 42
        session.environment_id = "env-1"
        session.current_phase = MagicMock(value="review")
        session.last_activity_at = None
        session.version = 1

        mock_repository.load_session_detail.return_value = session

        # Mock execution_snapshot with no blockers
        with patch("src.services.dataset_review.orchestrator.build_execution_snapshot") as mock_build_snapshot, \
             patch("src.services.dataset_review.orchestrator.SupersetCompilationAdapter") as mock_adapter_cls, \
             patch("src.services.dataset_review.orchestrator.build_launch_blockers") as mock_build_blockers:

            mock_snapshot = {
                "preview_blockers": [],
                "preview_fingerprint": "fp-1",
                "template_params": {"region": "US"},
                "effective_filters": [],
            }
            mock_build_snapshot.return_value = mock_snapshot

            mock_adapter = MagicMock()
            mock_adapter_cls.return_value = mock_adapter
            preview = MagicMock()
            preview.preview_id = "pv-1"
            preview.preview_status = PreviewStatus.READY
            preview.preview_fingerprint = "fp-1"
            mock_adapter.compile_preview.return_value = preview

            persisted_preview = MagicMock()
            persisted_preview.preview_id = "pv-1"
            persisted_preview.preview_status = PreviewStatus.READY
            mock_repository.save_preview.return_value = persisted_preview

            mock_build_blockers.return_value = []

            command = PreparePreviewCommand(
                session_id="sess-1",
                expected_version=1,
                user=MagicMock(id="user-1"),
            )

            result = orchestrator.prepare_launch_preview(command)

            assert result.session is not None
            assert result.session.session_id == "sess-1"

    def test_prepare_preview_session_not_found(self, orchestrator, mock_repository):
        """Missing session raises ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import PreparePreviewCommand

        mock_repository.load_session_detail.return_value = None

        command = PreparePreviewCommand(
            session_id="nonexistent",
            user=MagicMock(id="user-1"),
        )

        with pytest.raises(ValueError, match="Session not found"):
            orchestrator.prepare_launch_preview(command)

    def test_prepare_preview_no_dataset_id(self, orchestrator, mock_repository):
        """Missing dataset_id raises ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import PreparePreviewCommand

        session = MagicMock()
        session.session_id = "sess-1"
        session.user_id = "user-1"
        session.dataset_id = None
        session.environment_id = "env-1"
        mock_repository.load_session_detail.return_value = session

        command = PreparePreviewCommand(
            session_id="sess-1",
            user=MagicMock(id="user-1"),
        )

        with pytest.raises(ValueError, match="requires a resolved dataset_id"):
            orchestrator.prepare_launch_preview(command)

    def test_prepare_preview_env_not_found(self, orchestrator, mock_repository, mock_config_manager):
        """Missing environment raises ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import PreparePreviewCommand

        session = MagicMock()
        session.session_id = "sess-1"
        session.user_id = "user-1"
        session.dataset_id = 42
        session.environment_id = "nonexistent"
        session.current_phase = MagicMock(value="review")
        mock_repository.load_session_detail.return_value = session
        mock_config_manager.get_environment.return_value = None

        command = PreparePreviewCommand(
            session_id="sess-1",
            user=MagicMock(id="user-1"),
        )

        with pytest.raises(ValueError, match="Environment not found"):
            orchestrator.prepare_launch_preview(command)


# ── Tests for launch_dataset ───────────────────────────────

class TestLaunchDataset:
    """launch_dataset — dataset execution launch."""

    def test_launch_success(self, orchestrator, mock_repository, mock_config_manager):
        """Successful launch returns LaunchDatasetResult."""
        from src.services.dataset_review.orchestrator_pkg._commands import LaunchDatasetCommand
        from src.models.dataset_review import LaunchStatus

        session = MagicMock()
        session.session_id = "sess-1"
        session.user_id = "user-1"
        session.dataset_id = 42
        session.environment_id = "env-1"
        session.dataset_ref = "schema.table"
        session.current_phase = MagicMock(value="preview")
        session.last_activity_at = None
        session.version = 1
        mock_repository.load_session_detail.return_value = session

        with patch("src.services.dataset_review.orchestrator.build_execution_snapshot") as mock_build_snapshot, \
             patch("src.services.dataset_review.orchestrator.get_latest_preview") as mock_get_preview, \
             patch("src.services.dataset_review.orchestrator.build_launch_blockers") as mock_build_blockers, \
             patch("src.services.dataset_review.orchestrator.SupersetCompilationAdapter") as mock_adapter_cls:

            mock_snapshot = {
                "preview_blockers": [],
                "preview_fingerprint": "fp-1",
                "template_params": {"region": "US"},
                "effective_filters": [],
                "approved_mapping_ids": ["map-1"],
                "semantic_decision_refs": [],
                "open_warning_refs": [],
            }
            mock_build_snapshot.return_value = mock_snapshot

            preview = MagicMock()
            preview.preview_id = "pv-1"
            preview.compiled_sql = "SELECT * FROM table"
            preview.preview_fingerprint = "fp-1"
            mock_get_preview.return_value = preview

            mock_build_blockers.return_value = []

            mock_adapter = MagicMock()
            mock_adapter.create_sql_lab_session.return_value = "sql-lab-ref-1"
            mock_adapter_cls.return_value = mock_adapter

            run_context = MagicMock()
            run_context.run_context_id = "rc-1"
            run_context.launch_status = LaunchStatus.STARTED
            run_context.preview_id = "pv-1"
            run_context.sql_lab_session_ref = "sql-lab-ref-1"
            mock_repository.save_run_context.return_value = run_context

            command = LaunchDatasetCommand(
                session_id="sess-1",
                expected_version=1,
                user=MagicMock(id="user-1"),
            )

            result = orchestrator.launch_dataset(command)

            assert result.run_context is not None
            assert result.run_context.run_context_id == "rc-1"

    def test_launch_session_not_found(self, orchestrator, mock_repository):
        """Missing session raises ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import LaunchDatasetCommand

        mock_repository.load_session_detail.return_value = None

        command = LaunchDatasetCommand(
            session_id="nonexistent",
            user=MagicMock(id="user-1"),
        )

        with pytest.raises(ValueError, match="Session not found"):
            orchestrator.launch_dataset(command)

    def test_launch_blocked_by_validation(self, orchestrator, mock_repository):
        """Launch blockers raise ValueError."""
        from src.services.dataset_review.orchestrator_pkg._commands import LaunchDatasetCommand

        session = MagicMock()
        session.session_id = "sess-1"
        session.user_id = "user-1"
        session.dataset_id = 42
        session.environment_id = "env-1"
        session.current_phase = MagicMock(value="preview")
        mock_repository.load_session_detail.return_value = session

        with patch("src.services.dataset_review.orchestrator.build_execution_snapshot") as mock_build_snapshot, \
             patch("src.services.dataset_review.orchestrator.get_latest_preview") as mock_get_preview, \
             patch("src.services.dataset_review.orchestrator.build_launch_blockers") as mock_build_blockers:

            mock_build_snapshot.return_value = {"preview_blockers": []}
            preview = MagicMock()
            preview.preview_fingerprint = "fp-1"
            mock_get_preview.return_value = preview
            mock_build_blockers.return_value = ["finding:MISSING:blocking"]

            command = LaunchDatasetCommand(
                session_id="sess-1",
                user=MagicMock(id="user-1"),
            )

            with pytest.raises(ValueError, match="Launch blocked"):
                orchestrator.launch_dataset(command)

    def test_launch_sql_lab_failure(self, orchestrator, mock_repository, mock_config_manager):
        """SQL Lab failure sets launch status to FAILED without raising."""
        from src.services.dataset_review.orchestrator_pkg._commands import LaunchDatasetCommand
        from src.models.dataset_review import LaunchStatus

        session = MagicMock()
        session.session_id = "sess-1"
        session.user_id = "user-1"
        session.dataset_id = 42
        session.environment_id = "env-1"
        session.dataset_ref = "schema.table"
        session.current_phase = MagicMock(value="preview")
        session.last_activity_at = None
        session.version = 1
        mock_repository.load_session_detail.return_value = session

        with patch("src.services.dataset_review.orchestrator.build_execution_snapshot") as mock_build_snapshot, \
             patch("src.services.dataset_review.orchestrator.get_latest_preview") as mock_get_preview, \
             patch("src.services.dataset_review.orchestrator.build_launch_blockers") as mock_build_blockers, \
             patch("src.services.dataset_review.orchestrator.SupersetCompilationAdapter") as mock_adapter_cls:

            mock_build_snapshot.return_value = {
                "preview_blockers": [],
                "preview_fingerprint": "fp-1",
                "template_params": {},
                "effective_filters": [],
                "approved_mapping_ids": [],
                "semantic_decision_refs": [],
                "open_warning_refs": [],
            }
            preview = MagicMock()
            preview.preview_id = "pv-1"
            preview.compiled_sql = "SELECT * FROM table"
            preview.preview_fingerprint = "fp-1"
            mock_get_preview.return_value = preview
            mock_build_blockers.return_value = []

            mock_adapter = MagicMock()
            mock_adapter.create_sql_lab_session.side_effect = Exception("SQL Lab error")
            mock_adapter_cls.return_value = mock_adapter

            run_context = MagicMock()
            run_context.run_context_id = "rc-1"
            run_context.launch_status = LaunchStatus.FAILED
            run_context.preview_id = "pv-1"
            run_context.sql_lab_session_ref = "unavailable"
            mock_repository.save_run_context.return_value = run_context

            command = LaunchDatasetCommand(
                session_id="sess-1",
                user=MagicMock(id="user-1"),
            )

            result = orchestrator.launch_dataset(command)
            assert result.run_context.launch_status == LaunchStatus.FAILED


# ── Tests for _build_recovery_bootstrap ────────────────────

class TestBuildRecoveryBootstrap:
    """_build_recovery_bootstrap — recovery from Superset context."""

    @pytest.mark.asyncio
    async def test_build_recovery_bootstrap(self, orchestrator, mock_repository, mock_config_manager):
        """Recovery bootstrap builds filters, variables, and mappings."""
        session = MagicMock()
        session.session_id = "sess-1"
        session.dataset_id = 42

        parsed_context = MagicMock()
        parsed_context.dataset_payload = {"id": 42}
        parsed_context.unresolved_references = []

        environment = MagicMock()

        with patch("src.services.dataset_review.orchestrator.SupersetContextExtractor") as mock_extractor_cls:
            mock_extractor = MagicMock()
            mock_extractor.recover_imported_filters = AsyncMock(return_value=[
                {"filter_name": "region", "raw_value": "US"},
                {"filter_name": "date", "raw_value": "2024-01"}
            ])
            mock_extractor.discover_template_variables.return_value = [
                {"variable_name": "region", "expression_source": "table", "variable_kind": "native_filter", "is_required": True, "default_value": None}
            ]
            mock_extractor_cls.return_value = mock_extractor

            imported_filters, template_vars, mappings, findings = (
                await orchestrator._build_recovery_bootstrap(environment, session, parsed_context, [])
            )

        assert len(imported_filters) == 2
        assert imported_filters[0].filter_name == "region"
        assert len(template_vars) == 1
        assert template_vars[0].variable_name == "region"
        assert len(mappings) == 1
        assert mappings[0].mapping_method.value == "direct_match"

    @pytest.mark.asyncio
    async def test_recovery_bootstrap_handles_discovery_error(self, orchestrator, mock_repository):
        """Template variable discovery failure adds warning finding."""
        session = MagicMock()
        session.session_id = "sess-1"
        session.dataset_id = 42

        parsed_context = MagicMock()
        parsed_context.dataset_payload = None  # Will trigger fetch which fails
        parsed_context.unresolved_references = []

        environment = MagicMock()

        with patch("src.services.dataset_review.orchestrator.SupersetContextExtractor") as mock_extractor_cls:
            mock_extractor = MagicMock()
            mock_extractor.recover_imported_filters = AsyncMock(return_value=[])
            # Make the client call fail
            mock_client = MagicMock()
            mock_client.get_dataset_detail.side_effect = Exception("API error")
            mock_extractor.client = mock_client
            mock_extractor_cls.return_value = mock_extractor

            imported_filters, template_vars, mappings, findings = (
                await orchestrator._build_recovery_bootstrap(environment, session, parsed_context, [])
            )

        assert len(findings) == 1
        assert findings[0].code == "TEMPLATE_VARIABLE_DISCOVERY_FAILED"


# ── Tests for _enqueue_recovery_task ───────────────────────

class TestEnqueueRecoveryTask:
    """_enqueue_recovery_task — task linking."""

    def test_task_enqueued(self, orchestrator):
        """Task is created and ID returned when task_manager is present."""
        mock_task_manager = MagicMock()
        task = MagicMock()
        task.id = "task-1"
        mock_task_manager.create_task.return_value = task
        orchestrator.task_manager = mock_task_manager

        command = MagicMock()
        command.user.id = "user-1"
        session = MagicMock()
        session.session_id = "sess-1"
        session.environment_id = "env-1"
        session.source_kind = "dataset_selection"
        session.source_input = "42"
        session.dataset_ref = "dataset:42"
        session.dataset_id = 42
        session.dashboard_id = None

        task_id = orchestrator._enqueue_recovery_task(command, session, None)
        assert task_id == "task-1"

    def test_no_task_manager_returns_none(self, orchestrator):
        """No task_manager returns None."""
        orchestrator.task_manager = None
        command = MagicMock()
        session = MagicMock()
        result = orchestrator._enqueue_recovery_task(command, session, None)
        assert result is None

    def test_task_manager_missing_create_task(self, orchestrator):
        """Task manager without create_task returns None."""
        mock_task_manager = MagicMock(spec=[])
        orchestrator.task_manager = mock_task_manager

        command = MagicMock()
        session = MagicMock()
        result = orchestrator._enqueue_recovery_task(command, session, None)
        assert result is None

    def test_create_task_type_error_returns_none(self, orchestrator):
        """TypeError from create_task returns None."""
        mock_task_manager = MagicMock()
        mock_task_manager.create_task.side_effect = TypeError()
        orchestrator.task_manager = mock_task_manager

        command = MagicMock()
        command.user.id = "user-1"
        session = MagicMock()
        session.session_id = "sess-1"
        session.environment_id = "env-1"
        session.source_kind = "dataset_selection"
        session.source_input = "42"
        session.dataset_ref = "dataset:42"
        session.dataset_id = 42
        session.dashboard_id = None

        result = orchestrator._enqueue_recovery_task(command, session, None)
        assert result is None


# #endregion Test.DatasetReview.Orchestrator