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.
# @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.return_value = MagicMock(success=True, rows_affected=5, execution_time_ms=200)
                mock_exec_cls.return_value = mock_exec

                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.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.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