ss-tools/backend/tests/test_orchestrator_direct_db.py

# #region Test.Translate.OrchestratorSqlDirectDb [C:3] [TYPE Module] [SEMANTICS test, translate, orchestrator, direct-db]
# @BRIEF Pytest tests for SQLInsertService direct_db dispatch in orchestrator_sql.
# @RELATION BINDS_TO -> [SQLInsertService]
# @TEST_EDGE: missing_connection -> direct_db fails without connection_id
# @TEST_EDGE: invalid_insert_method -> Unsupported insert_method handled
# @TEST_EDGE: external_db_failure -> DB execution errors caught and reported
# @TEST_CONTRACT SQLInsertService dispatches to direct DB executor when job.insert_method == "direct_db".
# @TEST_CONTRACT Connection snapshot is stored on TranslationRun for direct DB inserts.

import pytest
from unittest.mock import MagicMock, patch, AsyncMock

from src.plugins.translate.orchestrator_sql import SQLInsertService


@pytest.fixture
def mock_db():
    db = MagicMock()
    db.query.return_value.options.return_value.filter.return_value.all.return_value = []
    return db


@pytest.fixture
def mock_config_manager():
    cm = MagicMock()
    cm.config.settings.connections = []
    return cm


@pytest.fixture
def mock_event_log():
    return MagicMock()


@pytest.fixture
def service(mock_db, mock_config_manager, mock_event_log):
    return SQLInsertService(mock_db, mock_config_manager, mock_event_log)


@pytest.fixture
def mock_job():
    job = MagicMock()
    job.id = "job-1"
    job.name = "Test Job"
    job.environment_id = "env-1"
    job.target_column = "translated_name"
    job.translation_column = "name"
    job.target_languages = ["en"]
    job.target_schema = None
    job.target_table = "products_i18n"
    job.target_key_cols = ["product_id"]
    job.upsert_strategy = "MERGE"
    job.database_dialect = "postgresql"
    job.target_dialect = "postgresql"
    job.insert_method = "sqllab"
    job.connection_id = None
    return job


@pytest.fixture
def mock_run():
    run = MagicMock()
    run.id = "run-1"
    run.connection_snapshot = None
    run.insert_method = None
    return run


class TestDirectDbDispatch:
    """SQLInsertService dispatches to DbExecutor (direct_db) or SupersetSqlLabExecutor (sqllab)."""

    @pytest.mark.asyncio
    async def test_direct_db_dispatch(self, service, mock_job, mock_run):
        """Given direct_db, dispatch to DbExecutor via ConnectionService, not SupersetSqlLabExecutor."""
        mock_job.insert_method = "direct_db"
        mock_job.connection_id = "conn-1"

        # Mock DbExecutor and ConnectionService at EXT boundaries
        mock_conn = MagicMock(name="Test PG", dialect="postgresql", host="x", port=5432, database="test_db")
        with patch('src.plugins.translate.orchestrator_sql.ConnectionService') as mock_cs_cls:
            mock_cs_cls.return_value.get_connection.return_value = mock_conn
            with patch('src.plugins.translate.orchestrator_sql.DbExecutor') as mock_exec_cls:
                mock_exec = MagicMock()
                mock_exec_cls.return_value = mock_exec
                mock_exec.execute_sql = AsyncMock(return_value=MagicMock(success=True, rows_affected=5, execution_time_ms=200))

                mock_record = MagicMock()
                mock_record.target_sql = "INSERT INTO ..."
                mock_record.status = "SUCCESS"
                service.db.query.return_value.options.return_value.filter.return_value.all.return_value = [mock_record]

                with patch('src.plugins.translate.orchestrator_sql.SQLGenerator.generate', return_value=("INSERT SQL", 5)):
                    result = await service.generate_and_insert_sql(mock_job, mock_run)

                    assert result["status"] == "success"
                    assert result["rows_affected"] == 5
                    mock_exec.execute_sql.assert_called_once_with("conn-1", "INSERT SQL")

    @pytest.mark.asyncio
    async def test_sqllab_dispatch_by_default(self, service, mock_job, mock_run):
        """Given default sqllab, dispatch to SupersetSqlLabExecutor."""
        mock_job.insert_method = "sqllab"
        mock_job.connection_id = None

        with patch('src.plugins.translate.orchestrator_sql.SupersetSqlLabExecutor') as mock_superset_cls:
            mock_superset = MagicMock()
            mock_superset_cls.return_value = mock_superset
            mock_superset.resolve_database_id = AsyncMock()
            mock_superset.execute_and_poll = AsyncMock(return_value={"status": "success"})

            mock_record = MagicMock()
            mock_record.target_sql = "INSERT INTO ..."
            mock_record.status = "SUCCESS"
            service.db.query.return_value.options.return_value.filter.return_value.all.return_value = [mock_record]

            with patch('src.plugins.translate.orchestrator_sql.SQLGenerator.generate', return_value=("INSERT SQL", 5)):
                result = await service.generate_and_insert_sql(mock_job, mock_run)

                assert result["status"] == "success"
                mock_superset.execute_and_poll.assert_called_once()


class TestDirectDbExecution:
    """SQLInsertService._execute_direct_db behavior."""

    @pytest.mark.asyncio
    async def test_connection_snapshot_stored(self, service, mock_job, mock_run):
        """Connection metadata is snapshotted on the run."""
        mock_job.insert_method = "direct_db"
        mock_job.connection_id = "conn-1"

        with patch('src.plugins.translate.orchestrator_sql.ConnectionService') as mock_cs_cls:
            mock_cs = MagicMock()
            mock_cs_cls.return_value = mock_cs
            mock_conn = MagicMock()
            mock_conn.name = "Test PG"
            mock_conn.dialect = "postgresql"
            mock_conn.host = "db.example.com"
            mock_conn.port = 5432
            mock_conn.database = "test_db"
            mock_cs.get_connection.return_value = mock_conn

            with patch('src.plugins.translate.orchestrator_sql.DbExecutor') as mock_exec_cls:
                mock_exec = MagicMock()
                mock_exec_cls.return_value = mock_exec
                mock_exec.execute_sql = AsyncMock(return_value=MagicMock(
                    success=True, rows_affected=10, execution_time_ms=200
                ))

                await service._execute_direct_db(mock_job, mock_run, "INSERT SQL")

                assert mock_run.connection_snapshot == {
                    "name": "Test PG",
                    "dialect": "postgresql",
                    "host": "db.example.com",
                    "port": 5432,
                    "database": "test_db",
                }
                assert mock_run.insert_method == "direct_db"

    @pytest.mark.asyncio
    async def test_connection_not_found(self, service, mock_job, mock_run):
        """Missing connection returns failure without crashing."""
        mock_job.insert_method = "direct_db"
        mock_job.connection_id = "nonexistent"

        with patch('src.plugins.translate.orchestrator_sql.ConnectionService') as mock_cs_cls:
            mock_cs = MagicMock()
            mock_cs_cls.return_value = mock_cs
            mock_cs.get_connection.return_value = None

            result = await service._execute_direct_db(mock_job, mock_run, "INSERT SQL")
            assert result["status"] == "failed"
            assert "not found" in result["error_message"]

    @pytest.mark.asyncio
    async def test_direct_db_execution_failure(self, service, mock_job, mock_run):
        """DB execution error returns failure details."""
        mock_job.insert_method = "direct_db"
        mock_job.connection_id = "conn-1"

        with patch('src.plugins.translate.orchestrator_sql.ConnectionService') as mock_cs_cls:
            mock_cs = MagicMock()
            mock_cs_cls.return_value = mock_cs
            mock_cs.get_connection.return_value = MagicMock(
                name="Test PG", dialect="postgresql", host="x", port=5432, database="test_db"
            )

            with patch('src.plugins.translate.orchestrator_sql.DbExecutor') as mock_exec_cls:
                mock_exec = MagicMock()
                mock_exec_cls.return_value = mock_exec
                mock_exec.execute_sql = AsyncMock(return_value=MagicMock(
                    success=False, rows_affected=0, error="Host unreachable", execution_time_ms=5000
                ))

                result = await service._execute_direct_db(mock_job, mock_run, "INSERT SQL")
                assert result["status"] == "failed"
                assert "Host unreachable" in result["error_message"]
                assert result["execution_time_ms"] == 5000


class TestGenerateAndInsertEdgeCases:
    """SQLInsertService.generate_and_insert_sql — edge cases (no records, empty columns, errors)."""

    @pytest.mark.asyncio
    async def test_no_successful_records_skipped(self, service, mock_job, mock_run):
        """When no SUCCESS records found, returns 'skipped' (covers lines 59-60)."""
        # mock_db.query already returns empty list from fixture
        result = await service.generate_and_insert_sql(mock_job, mock_run)
        assert result["status"] == "skipped"
        assert result["reason"] == "no_records"

    @pytest.mark.asyncio
    async def test_empty_columns_fallback(self, service, mock_job, mock_run):
        """When columns list is empty, fallback to [effective_target] (covers lines 70-71).

        NOTE: build_columns() always returns at least ['context', 'is_original'],
        so this path is currently dead code. The guard remains as a safety net
        in case build_columns is refactored. Keeping the test for completeness."""
        mock_job.target_key_cols = None
        mock_job.target_column = None
        mock_job.translation_column = None
        mock_job.target_language_column = None
        mock_job.target_source_column = None
        mock_job.target_source_language_column = None

        mock_record = MagicMock()
        mock_record.target_sql = "SELECT 1"
        mock_record.status = "SUCCESS"
        service.db.query.return_value.options.return_value.filter.return_value.all.return_value = [mock_record]

        with patch('src.plugins.translate.orchestrator_sql.SupersetSqlLabExecutor') as mock_superset_cls, \
             patch('src.plugins.translate.orchestrator_sql.SQLGenerator.generate', return_value=("INSERT SQL", 1)):
            mock_superset = MagicMock()
            mock_superset_cls.return_value = mock_superset
            mock_superset.resolve_database_id = AsyncMock()
            mock_superset.execute_and_poll = AsyncMock(return_value={"status": "success"})

            result = await service.generate_and_insert_sql(mock_job, mock_run)
            assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_sql_generation_value_error(self, service, mock_job, mock_run):
        """SQLGenerator raising ValueError returns 'failed' (covers lines 84-86)."""
        mock_record = MagicMock()
        mock_record.target_sql = "SELECT 1"
        mock_record.status = "SUCCESS"
        service.db.query.return_value.options.return_value.filter.return_value.all.return_value = [mock_record]

        with patch('src.plugins.translate.orchestrator_sql.SQLGenerator.generate', side_effect=ValueError("Bad upsert")):
            result = await service.generate_and_insert_sql(mock_job, mock_run)
            assert result["status"] == "failed"
            assert "Bad upsert" in result["error_message"]

    @pytest.mark.asyncio
    async def test_superset_execution_exception(self, service, mock_job, mock_run):
        """Superset execution raising Exception returns failed (covers lines 123-125)."""
        mock_record = MagicMock()
        mock_record.target_sql = "SELECT 1"
        mock_record.status = "SUCCESS"
        service.db.query.return_value.options.return_value.filter.return_value.all.return_value = [mock_record]

        with patch('src.plugins.translate.orchestrator_sql.SupersetSqlLabExecutor') as mock_superset_cls, \
             patch('src.plugins.translate.orchestrator_sql.SQLGenerator.generate', return_value=("INSERT SQL", 1)):
            mock_superset = MagicMock()
            mock_superset_cls.return_value = mock_superset
            mock_superset.resolve_database_id = AsyncMock()
            mock_superset.execute_and_poll = AsyncMock(side_effect=ConnectionError("Superset down"))

            result = await service.generate_and_insert_sql(mock_job, mock_run)
            assert result["status"] == "failed"
            assert "Superset down" in result["error_message"]


class TestResolveDialect:
    """SQLInsertService._resolve_dialect — dialect resolution edge cases."""

    @pytest.mark.asyncio
    async def test_direct_db_dialect(self, service, mock_job):
        """When insert_method == direct_db, dialect comes from ConnectionService (covers lines 176-180)."""
        mock_job.insert_method = "direct_db"
        mock_job.connection_id = "conn-1"

        mock_conn = MagicMock()
        mock_conn.dialect = "mysql"

        with patch('src.plugins.translate.orchestrator_sql.ConnectionService') as mock_cs_cls:
            mock_cs = MagicMock()
            mock_cs_cls.return_value = mock_cs
            mock_cs.get_connection.return_value = mock_conn

            dialect = await service._resolve_dialect(mock_job)
            assert dialect == "mysql"

    @pytest.mark.asyncio
    async def test_direct_db_no_connection_falls_through(self, service, mock_job):
        """When insert_method == direct_db but no connection, falls through to Superset (covers lines 178-180)."""
        mock_job.insert_method = "direct_db"
        mock_job.connection_id = "conn-missing"

        with patch('src.plugins.translate.orchestrator_sql.ConnectionService') as mock_cs_cls:
            mock_cs = MagicMock()
            mock_cs_cls.return_value = mock_cs
            mock_cs.get_connection.return_value = None

            with patch('src.plugins.translate.orchestrator_sql.SupersetSqlLabExecutor') as mock_superset_cls:
                mock_superset = MagicMock()
                mock_superset_cls.return_value = mock_superset
                mock_superset.resolve_database_id = AsyncMock()
                mock_superset.get_database_backend.return_value = "postgresql"

            dialect = await service._resolve_dialect(mock_job)
            assert dialect == "postgresql"
# #endregion Test.Translate.OrchestratorSqlDirectDb

    @pytest.mark.asyncio
    async def test_resolve_dialect_superset_exception(self, service, mock_job):
        """When Superset executor raises Exception, fallback to job dialect (covers lines 187-189)."""
        mock_job.insert_method = "sqllab"

        with patch('src.plugins.translate.orchestrator_sql.SupersetSqlLabExecutor') as mock_superset_cls:
            mock_superset = MagicMock()
            mock_superset_cls.return_value = mock_superset
            mock_superset.resolve_database_id = AsyncMock(side_effect=Exception("Superset unreachable"))

            dialect = await service._resolve_dialect(mock_job)
            # Falls back to job.database_dialect (postgresql from fixture)
            assert dialect == "postgresql"

    @pytest.mark.asyncio
    async def test_resolve_dialect_all_null_fallback(self, service, mock_job):
        """When all dialect sources are None, fallback to postgresql."""
        mock_job.database_dialect = None
        mock_job.target_dialect = None

        with patch('src.plugins.translate.orchestrator_sql.SupersetSqlLabExecutor') as mock_superset_cls:
            mock_superset = MagicMock()
            mock_superset_cls.return_value = mock_superset
            mock_superset.resolve_database_id = AsyncMock()
            mock_superset.get_database_backend.return_value = None

            dialect = await service._resolve_dialect(mock_job)
            assert dialect == "postgresql"