ss-tools/backend/tests/plugins/translate/test_batch_insert.py

# #region Test.BatchInsert [C:3] [TYPE Module] [SEMANTICS test,translate,batch,insert]
# @BRIEF Tests for _batch_insert.py — batch insert logic.
# @RELATION BINDS_TO -> [BatchInsertService]
# @TEST_EDGE: empty_batch -> Returns early
# @TEST_EDGE: no_columns_fallback -> Uses effective_target as single column
# @TEST_EDGE: backend_resolve_failure -> Returns early
# @TEST_EDGE: sql_generation_failure -> Returns early

import json
import sys
from datetime import UTC, datetime
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "src"))

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

from src.models.translate import (
    TranslationBatch,
    TranslationJob,
    TranslationRecord,
    TranslationLanguage,
    TranslationRun,
)
from src.plugins.translate._batch_insert import (
    _build_target_columns,
    _build_context_keys,
    _build_insert_rows,
    _generate_insert_sql,
    _execute_insert_sql,
    _resolve_insert_backend,
    _fetch_batch_records,
    insert_batch_to_target,
)
from .conftest import JOB_ID


class TestBuildTargetColumns:
    """_build_target_columns — build INSERT column list."""

    def test_basic_columns(self):
        """Basic columns with key_cols and effective_target."""
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = ["id"]
        job.target_language_column = "lang"
        job.target_source_column = "source"
        job.target_source_language_column = "src_lang"
        cols = _build_target_columns(job, "translated_text")
        assert "id" in cols
        assert "translated_text" in cols
        assert "lang" in cols
        assert "source" in cols
        assert "src_lang" in cols
        assert "context" in cols
        assert "is_original" in cols

    def test_deduplication(self):
        """Duplicate columns are removed."""
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = ["id", "id"]
        job.target_language_column = None
        job.target_source_column = None
        job.target_source_language_column = None
        cols = _build_target_columns(job, "id")
        # "id" appears only once despite being in key_cols and effective_target
        assert cols.count("id") == 1
        assert "context" in cols
        assert "is_original" in cols

    def test_empty_key_cols(self):
        """No key cols, no extra columns."""
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = None
        job.target_language_column = None
        job.target_source_column = None
        job.target_source_language_column = None
        cols = _build_target_columns(job, "text")
        assert cols == ["text", "context", "is_original"]

    def test_none_filtered(self):
        """None values in columns are filtered out."""
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = []
        job.target_language_column = None
        job.target_source_column = None
        job.target_source_language_column = None
        cols = _build_target_columns(job, None)
        assert "context" in cols
        assert "is_original" in cols
        assert None not in cols


class TestBuildContextKeys:
    """_build_context_keys — build context keys list."""

    def test_basic_context_keys(self):
        """Context from job.context_columns."""
        job = MagicMock(spec=TranslationJob)
        job.context_columns = ["col1", "col2"]
        job.translation_column = "translated_text"
        keys = _build_context_keys(job, "translated_text")
        assert keys == ["col1", "col2"]

    def test_translation_column_appended(self):
        """translation_column added when different from effective_target."""
        job = MagicMock(spec=TranslationJob)
        job.context_columns = ["col1"]
        job.translation_column = "orig_text"
        keys = _build_context_keys(job, "translated_text")
        assert "col1" in keys
        assert "orig_text" in keys

    def test_no_context_columns(self):
        """None context_columns returns empty list."""
        job = MagicMock(spec=TranslationJob)
        job.context_columns = None
        job.translation_column = "text"
        keys = _build_context_keys(job, "text")
        assert keys == []

    def test_no_duplicate_context(self):
        """translation_column not added if already in context_columns."""
        job = MagicMock(spec=TranslationJob)
        job.context_columns = ["text"]
        job.translation_column = "text"
        keys = _build_context_keys(job, "translated")
        # "text" already in context_columns, effective_target is different
        assert keys == ["text"]


class TestBuildInsertRows:
    """_build_insert_rows — build INSERT row dicts."""

    def _make_record(self, source_sql="hello", target_sql="bonjour",
                     source_data=None, languages=None, status="SUCCESS"):
        rec = MagicMock(spec=TranslationRecord)
        rec.source_sql = source_sql
        rec.target_sql = target_sql
        rec.source_data = source_data or {}
        rec.languages = languages or []
        rec.status = status
        return rec

    def test_basic_rows(self):
        """Basic case: record with languages."""
        lang = MagicMock(spec=TranslationLanguage)
        lang.language_code = "fr"
        lang.final_value = "bonjour"
        lang.translated_value = "bonjour"
        lang.source_language_detected = "en"

        rec = self._make_record(
            source_sql="hello", target_sql="bonjour",
            source_data={"id": 1}, languages=[lang],
        )
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = ["id"]
        job.target_source_column = "source"
        job.target_source_language_column = "src_lang"
        job.target_language_column = "lang"
        job.context_columns = []

        rows = _build_insert_rows([rec], job, "translated_text", "en", [])
        # Should have 2 rows: original + translation
        assert len(rows) == 2
        original = rows[0]
        translation = rows[1]
        assert original["is_original"] == 1
        assert translation["is_original"] == 0
        assert translation["lang"] == "fr"
        assert translation["translated_text"] == "bonjour"
        assert original["translated_text"] == "hello"

    def test_no_languages_fallback(self):
        """Record with no languages -> fallback row with primary_language."""
        rec = self._make_record(
            source_sql="hello", target_sql="bonjour",
            source_data={"id": 1}, languages=[],
        )
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = ["id"]
        job.target_source_column = None
        job.target_source_language_column = None
        job.target_language_column = "lang"
        job.context_columns = []

        rows = _build_insert_rows([rec], job, "translated_text", "fr", [])
        assert len(rows) == 2
        assert rows[1]["lang"] == "fr"
        assert rows[1]["is_original"] == 0
        assert rows[1]["translated_text"] == "bonjour"

    def test_skip_same_language(self):
        """Skip translation row with same language as detected source."""
        lang = MagicMock(spec=TranslationLanguage)
        lang.language_code = "en"
        lang.final_value = "hello"
        lang.translated_value = "hello"
        lang.source_language_detected = "en"

        rec = self._make_record(
            source_sql="hello", target_sql="bonjour",
            source_data={"id": 1}, languages=[lang],
        )
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = []
        job.target_source_column = None
        job.target_source_language_column = None
        job.target_language_column = None
        job.context_columns = []

        rows = _build_insert_rows([rec], job, "translated_text", "en", [])
        # Only the original row (fr translation has same code as detected "en")
        # Wait: lang.language_code == "en" but detected_src_lang from rec.languages[0].source_language_detected = "en"
        # So the skip condition fires and only original is kept
        assert len(rows) == 1

    def test_no_source_language_detected(self):
        """When rec.languages has source_language_detected missing."""
        lang = MagicMock(spec=TranslationLanguage)
        lang.language_code = "fr"
        lang.final_value = "bonjour"
        lang.translated_value = "bonjour"
        lang.source_language_detected = None

        rec = self._make_record(
            source_sql="hello", target_sql="bonjour",
            source_data={}, languages=[lang],
        )
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = []
        job.target_source_column = None
        job.target_source_language_column = None
        job.target_language_column = None
        job.context_columns = []

        rows = _build_insert_rows([rec], job, "translated_text", "en", [])
        assert len(rows) == 2

    def test_context_data_included(self):
        """Context keys are collected."""
        lang = MagicMock(spec=TranslationLanguage)
        lang.language_code = "fr"
        lang.final_value = "bonjour"
        lang.translated_value = "bonjour"
        lang.source_language_detected = "en"

        rec = self._make_record(
            source_sql="hello", target_sql="bonjour",
            source_data={"id": 1, "desc": "test"}, languages=[lang],
        )
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = ["id"]
        job.target_source_column = None
        job.target_source_language_column = None
        job.target_language_column = None
        job.context_columns = []

        rows = _build_insert_rows([rec], job, "translated_text", "en", [])
        context = json.loads(rows[0]["context"])
        assert context == {}  # no context_keys

    def test_with_context_keys(self):
        """Context keys are included in context JSON."""
        lang = MagicMock(spec=TranslationLanguage)
        lang.language_code = "fr"
        lang.final_value = "bonjour"
        lang.translated_value = "bonjour"
        lang.source_language_detected = "en"

        rec = self._make_record(
            source_sql="hello", target_sql="bonjour",
            source_data={"id": 1, "desc": "test"}, languages=[lang],
        )
        job = MagicMock(spec=TranslationJob)
        job.target_key_cols = ["id"]
        job.target_source_column = None
        job.target_source_language_column = None
        job.target_language_column = None
        job.context_columns = []

        rows = _build_insert_rows([rec], job, "translated_text", "en", ["desc"])
        context = json.loads(rows[0]["context"])
        assert context.get("desc") == "test"


class TestFetchBatchRecords:
    """_fetch_batch_records — DB query."""

    def _create_run_and_batch(self, db_session):
        """Helper to create a TranslationRun and TranslationBatch for FK compliance."""
        from src.models.translate import TranslationRun, TranslationBatch
        from datetime import datetime, UTC
        run = TranslationRun(id="test-run-id", job_id=JOB_ID, status="RUNNING",
                             started_at=datetime.now(UTC))
        db_session.add(run)
        db_session.flush()
        batch = TranslationBatch(id="batch-1", run_id="test-run-id",
                                 batch_index=0, status="PENDING")
        db_session.add(batch)
        db_session.commit()

    def test_returns_records(self, db_session):
        """Returns records with SUCCESS status and target_sql."""
        self._create_run_and_batch(db_session)
        from src.models.translate import TranslationRecord, TranslationLanguage
        rec = TranslationRecord(
            id="rec-1", batch_id="batch-1", run_id="test-run-id",
            source_sql="hello", target_sql="bonjour", status="SUCCESS",
        )
        db_session.add(rec)
        db_session.commit()

        results = _fetch_batch_records(db_session, "batch-1")
        assert len(results) == 1
        assert results[0].id == "rec-1"

    def test_skips_non_success(self, db_session):
        """Records without SUCCESS status are excluded."""
        self._create_run_and_batch(db_session)
        from src.models.translate import TranslationRecord
        rec = TranslationRecord(
            id="rec-2", batch_id="batch-1", run_id="test-run-id",
            source_sql="hello", target_sql="bonjour", status="FAILED",
        )
        db_session.add(rec)
        db_session.commit()

        results = _fetch_batch_records(db_session, "batch-1")
        assert len(results) == 0

    def test_skips_no_target_sql(self, db_session):
        """Records with null target_sql are excluded."""
        self._create_run_and_batch(db_session)
        from src.models.translate import TranslationRecord
        rec = TranslationRecord(
            id="rec-3", batch_id="batch-1", run_id="test-run-id",
            source_sql="hello", target_sql=None, status="SUCCESS",
        )
        db_session.add(rec)
        db_session.commit()

        results = _fetch_batch_records(db_session, "batch-1")
        assert len(results) == 0

    def test_returns_empty_for_missing_batch(self, db_session):
        """Non-existent batch returns empty list."""
        results = _fetch_batch_records(db_session, "nonexistent")
        assert results == []


class TestGenerateInsertSql:
    """_generate_insert_sql — SQL generation."""

    def test_generates_sql(self):
        """Returns SQL and row count."""
        job = MagicMock(spec=TranslationJob)
        job.target_schema = "public"
        job.target_table = "translations"
        job.target_key_cols = None
        job.upsert_strategy = "MERGE"

        rows = [{"col1": "val1", "is_original": 1}]
        sql, count = _generate_insert_sql("postgresql", job, ["col1", "is_original"], rows, "batch-1")
        assert sql is not None
        assert "INSERT" in sql.upper()
        assert count > 0

    def test_value_error_returns_none(self):
        """ValueError in SQL generation returns (None, 0)."""
        job = MagicMock(spec=TranslationJob)
        job.target_schema = "public"
        job.target_table = None  # Will cause error
        job.target_key_cols = None
        job.upsert_strategy = "MERGE"

        with patch("src.plugins.translate._batch_insert.SQLGenerator.generate",
                   side_effect=ValueError("test error")):
            sql, count = _generate_insert_sql("postgresql", job, ["col"], [{"col": "v"}], "batch-1")
            assert sql is None
            assert count == 0


class TestExecuteInsertSql:
    """_execute_insert_sql — execute via Superset."""

    @pytest.mark.asyncio
    async def test_success(self):
        """Successful execution logs reason."""
        executor = AsyncMock()
        executor.execute_and_poll.return_value = {"status": "success"}
        await _execute_insert_sql(executor, "INSERT ...", "batch-1", 5)
        executor.execute_and_poll.assert_called_once_with(sql="INSERT ...", max_polls=30, poll_interval_seconds=2.0)

    @pytest.mark.asyncio
    async def test_exception_caught(self):
        """Exception in execute_and_poll is caught gracefully."""
        executor = AsyncMock()
        executor.execute_and_poll.side_effect = Exception("API error")
        await _execute_insert_sql(executor, "INSERT ...", "batch-1", 5)
        executor.execute_and_poll.assert_called_once()

    @pytest.mark.asyncio
    async def test_direct_db_success(self):
        """DbExecutor path: successful execution."""
        from src.core.db_executor import DbExecutor, DbExecutionResult
        executor = MagicMock(spec=DbExecutor)
        executor.execute_sql = AsyncMock(return_value=DbExecutionResult(
            success=True, rows_affected=5, execution_time_ms=12.0
        ))
        await _execute_insert_sql(executor, "INSERT ...", "batch-1", 5, connection_id="conn-1")
        executor.execute_sql.assert_called_once_with("conn-1", "INSERT ...")

    @pytest.mark.asyncio
    async def test_direct_db_failure(self):
        """DbExecutor path: failed execution with error."""
        from src.core.db_executor import DbExecutor, DbExecutionResult
        executor = MagicMock(spec=DbExecutor)
        executor.execute_sql = AsyncMock(return_value=DbExecutionResult(
            success=False, error="Connection refused"
        ))
        await _execute_insert_sql(executor, "INSERT ...", "batch-1", 5, connection_id="conn-1")
        executor.execute_sql.assert_called_once_with("conn-1", "INSERT ...")


class TestResolveInsertBackend:
    """_resolve_insert_backend — resolve backend dialect and executor."""

    @pytest.mark.asyncio
    async def test_success(self):
        """Returns dialect and executor."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.environment_id = "env-1"
        job.source_dialect = None
        job.target_database_id = None
        job.database_dialect = None
        job.target_dialect = None

        with patch("src.plugins.translate._batch_insert.SupersetSqlLabExecutor") as mock_exec_cls:
            mock_exec = MagicMock()
            mock_exec.get_database_backend.return_value = "postgresql"
            mock_exec.resolve_database_id = AsyncMock()
            mock_exec_cls.return_value = mock_exec

            dialect, executor, connection_id = await _resolve_insert_backend(config_manager, job, "batch-1")
            assert dialect == "postgresql"
            assert executor is not None
            assert connection_id is None  # sqllab path has no connection_id

    @pytest.mark.asyncio
    async def test_failure_returns_none(self):
        """Exception in resolve returns (None, None, None)."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.environment_id = None
        job.source_dialect = None

        with patch("src.plugins.translate._batch_insert.SupersetSqlLabExecutor",
                   side_effect=Exception("Failed")):
            dialect, executor, connection_id = await _resolve_insert_backend(config_manager, job, "batch-1")
            assert dialect is None
            assert executor is None
            assert connection_id is None

    @pytest.mark.asyncio
    async def test_fallback_dialect(self):
        """Falls back to job.database_dialect or target_dialect."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.environment_id = None
        job.source_dialect = None
        job.target_database_id = None
        job.database_dialect = "mysql"
        job.target_dialect = None

        with patch("src.plugins.translate._batch_insert.SupersetSqlLabExecutor") as mock_exec_cls:
            mock_exec = MagicMock()
            mock_exec.get_database_backend.return_value = None
            mock_exec.resolve_database_id = AsyncMock()
            mock_exec_cls.return_value = mock_exec

            dialect, _, connection_id = await _resolve_insert_backend(config_manager, job, "batch-1")
            assert dialect == "mysql"
            assert connection_id is None

    @pytest.mark.asyncio
    async def test_direct_db_path(self):
        """direct_db path returns DbExecutor with connection_id."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.insert_method = "direct_db"
        job.connection_id = "clickhouse-conn-1"

        with patch("src.plugins.translate._batch_insert.ConnectionService") as mock_cs_cls:
            mock_cs = MagicMock()
            mock_conn = MagicMock()
            mock_conn.dialect = "clickhouse"
            mock_conn.name = "Test ClickHouse"
            mock_cs.get_connection.return_value = mock_conn
            mock_cs_cls.return_value = mock_cs

            dialect, executor, connection_id = await _resolve_insert_backend(config_manager, job, "batch-1")
            assert dialect == "clickhouse"
            assert connection_id == "clickhouse-conn-1"
            # Should return DbExecutor, not SupersetSqlLabExecutor
            from src.core.db_executor import DbExecutor
            assert isinstance(executor, DbExecutor)

    @pytest.mark.asyncio
    async def test_direct_db_missing_conn_raises(self):
        """When direct_db connection not found, raises ValueError (no fallback)."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.insert_method = "direct_db"
        job.connection_id = "missing-conn"

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

            with pytest.raises(ValueError, match="connection_id='missing-conn' not found"):
                await _resolve_insert_backend(config_manager, job, "batch-1")

    @pytest.mark.asyncio
    async def test_direct_db_no_connection_id_raises(self):
        """When direct_db is set but connection_id is None, raises ValueError (no fallback)."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.insert_method = "direct_db"
        job.connection_id = None

        with pytest.raises(ValueError, match="connection_id is missing"):
            await _resolve_insert_backend(config_manager, job, "batch-1")


class TestInsertBatchToTarget:
    """insert_batch_to_target — main integration function."""

    @pytest.mark.asyncio
    async def test_empty_batch_returns_early(self, db_session):
        """Empty batch returns without error."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.target_column = None
        job.translation_column = "text"
        job.target_languages = None

        with patch("src.plugins.translate._batch_insert._fetch_batch_records", return_value=[]):
            result = await insert_batch_to_target(db_session, config_manager, job, "batch-1", "run-1")
            assert result is None

    @pytest.mark.asyncio
    async def test_no_columns_fallback(self, db_session):
        """When columns empty, falls back to single column."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.target_column = None
        job.translation_column = None
        job.target_languages = None
        job.target_schema = None
        job.target_table = None
        job.target_key_cols = None
        job.upsert_strategy = "MERGE"
        job.environment_id = None
        job.source_dialect = None
        job.target_database_id = None
        job.database_dialect = "postgresql"
        job.target_dialect = None

        rec = MagicMock(spec=TranslationRecord)
        rec.target_sql = "bonjour"

        with patch("src.plugins.translate._batch_insert._fetch_batch_records", return_value=[rec]):
            with patch("src.plugins.translate._batch_insert._resolve_insert_backend",
                       return_value=("postgresql", AsyncMock(), None)):
                with patch("src.plugins.translate._batch_insert.SQLGenerator.generate_batch",
                           return_value=[("INSERT INTO ...", 1)]):
                    with patch("src.plugins.translate._batch_insert._execute_insert_sql",
                               AsyncMock()):
                        result = await insert_batch_to_target(
                            db_session, config_manager, job, "batch-1", "run-1"
                        )
                        assert result is None

    @pytest.mark.asyncio
    async def test_backend_resolve_none_returns_early(self, db_session):
        """When _resolve_insert_backend returns None, returns early."""
        config_manager = MagicMock()
        job = MagicMock(spec=TranslationJob)
        job.target_column = "text"
        job.translation_column = None
        job.target_languages = []

        rec = MagicMock(spec=TranslationRecord)
        rec.target_sql = "bonjour"
        rec.source_sql = "hello"
        rec.source_data = {}
        rec.languages = []

        with patch("src.plugins.translate._batch_insert._fetch_batch_records", return_value=[rec]):
            with patch("src.plugins.translate._batch_insert._build_target_columns", return_value=["text"]):
                with patch("src.plugins.translate._batch_insert._build_context_keys", return_value=[]):
                    with patch("src.plugins.translate._batch_insert._build_insert_rows",
                               return_value=[{"text": "bonjour"}]):
                        with patch("src.plugins.translate._batch_insert._resolve_insert_backend",
                                   return_value=(None, None, None)):
                            result = await insert_batch_to_target(
                                db_session, config_manager, job, "batch-1", "run-1"
                            )
                            assert result is None
# #endregion Test.BatchInsert