ss-tools/backend/src/plugins/translate/__tests__/test_preview.py

# region TranslationPreviewTests [TYPE Module]
# @SEMANTICS: test, translate, preview, session
# @PURPOSE: Test TranslationPreview service: preview, approve/edit/reject state transitions, cost estimation, acceptance gate.
# @LAYER: Test
# @RELATION: BINDS_TO -> [TranslationPreview:Module]

import json
from datetime import UTC, datetime, timedelta
from unittest.mock import MagicMock, patch

import pytest

from src.models.translate import (
    TranslationJob,
    TranslationPreviewRecord,
    TranslationPreviewSession,
)
from src.plugins.translate.preview import TokenEstimator, TranslationPreview


# region _make_mock_job [TYPE Function]
# @PURPOSE: Create a mock TranslationJob with test config.
def _make_mock_job(**overrides):
    job = MagicMock(spec=TranslationJob)
    job.id = overrides.get("id", "job-123")
    job.name = overrides.get("name", "Test Job")
    job.source_dialect = overrides.get("source_dialect", "postgresql")
    job.target_dialect = overrides.get("target_dialect", "clickhouse")
    job.source_datasource_id = overrides.get("source_datasource_id", "42")
    job.source_table = overrides.get("source_table", "source_table")
    job.translation_column = overrides.get("translation_column", "title")
    job.context_columns = overrides.get("context_columns", ["category", "description"])
    job.target_language = overrides.get("target_language", "ru")
    job.provider_id = overrides.get("provider_id", "provider-1")
    job.batch_size = overrides.get("batch_size", 50)
    job.upsert_strategy = overrides.get("upsert_strategy", "MERGE")
    job.status = overrides.get("status", "DRAFT")
    job.target_column = None
    return job
# endregion _make_mock_job


# region _make_mock_provider [TYPE Function]
# @PURPOSE: Create a mock LLMProvider.
def _make_mock_provider(**overrides):
    provider = MagicMock()
    provider.id = overrides.get("id", "provider-1")
    provider.provider_type = overrides.get("provider_type", "openai")
    provider.base_url = overrides.get("base_url", "https://api.openai.com/v1")
    provider.default_model = overrides.get("default_model", "gpt-4o-mini")
    provider.api_key = "encrypted-key-123"
    return provider
# endregion _make_mock_provider


# region TestTranslationPreview [TYPE Class]
# @PURPOSE: Test suite for TranslationPreview service.
class TestTranslationPreview:

    # region test_preview_valid_job [TYPE Function]
    # @PURPOSE: Verify preview creates session and records successfully.
    def test_preview_valid_job(self):
        """Preview with valid job should create session and return records."""
        db = MagicMock()
        config_manager = MagicMock()

        # Mock job query
        job = _make_mock_job()
        db.query.return_value.filter.return_value.first.return_value = job

        # Mock environments
        env = MagicMock()
        env.id = "postgresql"
        env.url = "https://superset.example.com"
        env.username = "admin"
        env.password = "admin"
        env.verify_ssl = True
        env.timeout = 30
        config_manager.get_environments.return_value = [env]

        # Mock SupersetClient
        with patch("src.plugins.translate.preview.SupersetClient") as MockClient:
            mock_client = MagicMock()
            MockClient.return_value = mock_client

            # Mock dataset_detail
            mock_client.get_dataset_detail.return_value = {
                "table_name": "test_table",
                "columns": [
                    {"column_name": "title", "type": "VARCHAR"},
                    {"column_name": "category", "type": "VARCHAR"},
                    {"column_name": "description", "type": "TEXT"},
                ],
            }

            # Mock build_dataset_preview_query_context
            mock_client.build_dataset_preview_query_context.return_value = {
                "datasource": {"id": 42, "type": "table"},
                "queries": [{
                    "columns": [],
                    "metrics": ["count"],
                    "row_limit": 10,
                }],
                "form_data": {},
                "result_format": "json",
                "result_type": "query",
                "force": True,
            }

            # Mock network request to Superset
            mock_client.network.request.return_value = {
                "result": [{
                    "status": "success",
                    "data": [
                        {"title": "Hello World", "category": "greeting", "description": "A friendly message"},
                        {"title": "Goodbye", "category": "farewell", "description": "A parting message"},
                    ],
                }]
            }

            # Mock LLM provider
            with patch("src.plugins.translate.preview.LLMProviderService") as MockProviderService:
                mock_provider_svc = MagicMock()
                MockProviderService.return_value = mock_provider_svc
                mock_provider = _make_mock_provider()
                mock_provider_svc.get_provider.return_value = mock_provider
                mock_provider_svc.get_decrypted_api_key.return_value = "sk-test-key"

                # Mock _call_openai_compatible to return structured JSON
                with patch.object(TranslationPreview, "_call_openai_compatible") as mock_llm:
                    mock_llm.return_value = json.dumps({
                        "rows": [
                            {"row_id": "0", "translation": "Привет мир"},
                            {"row_id": "1", "translation": "Прощай"},
                        ]
                    })

                    # Mock DictionaryManager.filter_for_batch
                    with patch("src.plugins.translate.preview.DictionaryManager.filter_for_batch") as mock_filter:
                        mock_filter.return_value = []

                        preview_service = TranslationPreview(db, config_manager, "test-user")
                        result = preview_service.preview_rows(job_id="job-123", sample_size=10)

                        # Verify result structure
                        assert result["job_id"] == "job-123"
                        assert result["status"] == "ACTIVE"
                        assert len(result["records"]) == 2
                        assert result["records"][0]["target_sql"] == "Привет мир"
                        assert result["records"][1]["target_sql"] == "Прощай"

                        # Verify cost estimate
                        assert result["cost_estimate"]["sample_size"] == 2
                        assert result["cost_estimate"]["sample_total_tokens"] > 0
                        assert result["cost_estimate"]["sample_cost"] >= 0

                        # Verify config and dict hashes
                        assert result["config_hash"] is not None
                        assert result["dict_snapshot_hash"] is not None

                        # Verify session was added to DB
                        assert db.add.called
                        assert db.commit.called
    # endregion test_preview_valid_job

    # region test_preview_with_dictionary [TYPE Function]
    # @PURPOSE: Verify glossary entries from dictionary are included in LLM prompt.
    def test_preview_with_dictionary(self):
        """Preview should include dictionary glossary in the prompt sent to LLM."""
        db = MagicMock()
        config_manager = MagicMock()

        job = _make_mock_job()
        db.query.return_value.filter.return_value.first.return_value = job

        env = MagicMock()
        env.id = "postgresql"
        config_manager.get_environments.return_value = [env]

        with patch("src.plugins.translate.preview.SupersetClient") as MockClient:
            mock_client = MagicMock()
            MockClient.return_value = mock_client
            mock_client.get_dataset_detail.return_value = {"table_name": "test", "columns": []}
            mock_client.build_dataset_preview_query_context.return_value = {
                "datasource": {"id": 42, "type": "table"},
                "queries": [{"columns": [], "metrics": ["count"], "row_limit": 10}],
                "form_data": {},
                "result_format": "json",
                "result_type": "query",
                "force": True,
            }
            mock_client.network.request.return_value = {
                "result": [{"status": "success", "data": [{"title": "Hello World", "category": "greeting"}]}]
            }

            with patch("src.plugins.translate.preview.LLMProviderService") as MockProviderService:
                mock_provider_svc = MagicMock()
                MockProviderService.return_value = mock_provider_svc
                mock_provider_svc.get_provider.return_value = _make_mock_provider()
                mock_provider_svc.get_decrypted_api_key.return_value = "sk-test-key"

                with patch.object(TranslationPreview, "_call_openai_compatible") as mock_llm:
                    mock_llm.return_value = json.dumps({"rows": [{"row_id": "0", "translation": "Привет мир"}]})

                    # Mock dictionary filter to return glossary entries
                    with patch("src.plugins.translate.preview.DictionaryManager.filter_for_batch") as mock_filter:
                        mock_filter.return_value = [
                            {
                                "source_term": "Hello",
                                "target_term": "Здравствуйте",
                                "dictionary_id": "dict-1",
                                "dictionary_name": "Test Dict",
                                "context_notes": "Formal greeting",
                            }
                        ]

                        preview_service = TranslationPreview(db, config_manager, "test-user")
                        result = preview_service.preview_rows(job_id="job-123", sample_size=5)

                        # Verify LLM was called
                        mock_llm.assert_called_once()
                        assert len(result["records"]) == 1
    # endregion test_preview_with_dictionary

    # region test_preview_row_approve [TYPE Function]
    # @PURPOSE: Verify approving a preview row changes its status.
    def test_preview_row_approve(self):
        """Approve action should set record status to APPROVED."""
        db = MagicMock()
        config_manager = MagicMock()

        # Mock active session
        session = MagicMock(spec=TranslationPreviewSession)
        session.id = "session-1"
        session.job_id = "job-123"
        session.status = "ACTIVE"
        session.created_at = datetime.now(UTC)

        # Mock record
        record = MagicMock(spec=TranslationPreviewRecord)
        record.id = "record-1"
        record.session_id = "session-1"
        record.status = "PENDING"
        record.target_sql = "test translation"
        record.feedback = None

        # Setup: db.query().filter().order_by().first() -> session (session lookup)
        db.query.return_value.filter.return_value.order_by.return_value.first.return_value = session
        # Setup: db.query().filter().first() -> record (record lookup)
        db.query.return_value.filter.return_value.first.return_value = record

        preview_service = TranslationPreview(db, config_manager, "test-user")
        result = preview_service.update_preview_row(
            job_id="job-123",
            row_id="record-1",
            action="approve",
        )

        assert record.status == "APPROVED"
        assert result["status"] == "APPROVED"
        assert db.commit.called
    # endregion test_preview_row_approve

    # region test_preview_row_reject [TYPE Function]
    # @PURPOSE: Verify rejecting a preview row changes its status.
    def test_preview_row_reject(self):
        """Reject action should set record status to REJECTED."""
        db = MagicMock()
        config_manager = MagicMock()

        session = MagicMock(spec=TranslationPreviewSession)
        session.id = "session-1"
        session.job_id = "job-123"
        session.status = "ACTIVE"

        record = MagicMock(spec=TranslationPreviewRecord)
        record.id = "record-1"
        record.session_id = "session-1"
        record.status = "PENDING"
        record.target_sql = "test"
        record.feedback = None

        db.query.return_value.filter.return_value.order_by.return_value.first.return_value = session
        db.query.return_value.filter.return_value.first.return_value = record

        preview_service = TranslationPreview(db, config_manager, "test-user")
        result = preview_service.update_preview_row(
            job_id="job-123",
            row_id="record-1",
            action="reject",
        )

        assert record.status == "REJECTED"
        assert result["status"] == "REJECTED"
    # endregion test_preview_row_reject

    # region test_preview_row_edit [TYPE Function]
    # @PURPOSE: Verify editing a preview row updates its translation and status.
    def test_preview_row_edit(self):
        """Edit action should update translation and set status to APPROVED."""
        db = MagicMock()
        config_manager = MagicMock()

        session = MagicMock(spec=TranslationPreviewSession)
        session.id = "session-1"
        session.job_id = "job-123"
        session.status = "ACTIVE"

        record = MagicMock(spec=TranslationPreviewRecord)
        record.id = "record-1"
        record.session_id = "session-1"
        record.status = "PENDING"
        record.target_sql = "old translation"
        record.feedback = None

        db.query.return_value.filter.return_value.order_by.return_value.first.return_value = session
        db.query.return_value.filter.return_value.first.return_value = record

        preview_service = TranslationPreview(db, config_manager, "test-user")
        result = preview_service.update_preview_row(
            job_id="job-123",
            row_id="record-1",
            action="edit",
            translation="edited translation",
        )

        assert record.target_sql == "edited translation"
        assert record.status == "APPROVED"
        assert result["status"] == "APPROVED"
        assert result["target_sql"] == "edited translation"
    # endregion test_preview_row_edit

    # region test_preview_row_invalid_action [TYPE Function]
    # @PURPOSE: Verify invalid action raises ValueError.
    def test_preview_row_invalid_action(self):
        """Invalid action should raise ValueError."""
        db = MagicMock()
        config_manager = MagicMock()

        session = MagicMock(spec=TranslationPreviewSession)
        session.id = "session-1"
        session.job_id = "job-123"
        session.status = "ACTIVE"

        record = MagicMock(spec=TranslationPreviewRecord)
        record.id = "record-1"
        record.session_id = "session-1"
        record.status = "PENDING"

        db.query.return_value.filter.return_value.order_by.return_value.first.return_value = session
        db.query.return_value.filter.return_value.first.return_value = record

        preview_service = TranslationPreview(db, config_manager, "test-user")
        with pytest.raises(ValueError, match="Invalid action"):
            preview_service.update_preview_row(
                job_id="job-123",
                row_id="record-1",
                action="invalid_action",
            )
    # endregion test_preview_row_invalid_action

    # region test_preview_accept_session [TYPE Function]
    # @PURPOSE: Verify accepting a preview session gates execution.
    def test_preview_accept_session(self):
        """Accept should set session status to APPLIED and return records."""
        db = MagicMock()
        config_manager = MagicMock()

        session = MagicMock(spec=TranslationPreviewSession)
        session.id = "session-1"
        session.job_id = "job-123"
        session.status = "ACTIVE"
        session.created_by = "test-user"
        session.created_at = datetime.now(UTC)
        session.expires_at = datetime.now(UTC) + timedelta(hours=24)

        record = MagicMock(spec=TranslationPreviewRecord)
        record.id = "record-1"
        record.session_id = "session-1"
        record.source_sql = "Hello"
        record.target_sql = "Привет"
        record.status = "APPROVED"
        record.feedback = None

        # Session lookup
        db.query.return_value.filter.return_value.order_by.return_value.first.return_value = session
        # Records lookup
        db.query.return_value.filter.return_value.all.return_value = [record]

        preview_service = TranslationPreview(db, config_manager, "test-user")
        result = preview_service.accept_preview_session(job_id="job-123")

        assert session.status == "APPLIED"
        assert result["status"] == "APPLIED"
        assert len(result["records"]) == 1
        assert result["records"][0]["status"] == "APPROVED"
        assert db.commit.called
    # endregion test_preview_accept_session

    # region test_preview_no_active_session [TYPE Function]
    # @PURPOSE: Verify accept raises error when no active session.
    def test_preview_no_active_session(self):
        """Accept without active session should raise ValueError."""
        db = MagicMock()
        config_manager = MagicMock()

        db.query.return_value.filter.return_value.order_by.return_value.first.return_value = None

        preview_service = TranslationPreview(db, config_manager, "test-user")
        with pytest.raises(ValueError, match="No active preview session"):
            preview_service.accept_preview_session(job_id="job-123")
    # endregion test_preview_no_active_session

    # region test_cost_estimation [TYPE Function]
    # @PURPOSE: Verify token and cost estimation methods.
    def test_cost_estimation(self):
        """Token and cost estimation should return reasonable values."""
        prompt = "Translate this text from English to Russian. " * 100
        tokens = TokenEstimator.estimate_prompt_tokens(prompt)
        assert tokens > 0
        assert tokens < len(prompt)  # tokens typically fewer than chars

        # Single language
        output_tokens = TokenEstimator.estimate_output_tokens(10, num_languages=1)
        assert output_tokens == 600  # 10 * 1 * 50 * 1.2

        # Multi-language
        output_tokens_multi = TokenEstimator.estimate_output_tokens(10, num_languages=3)
        assert output_tokens_multi == 1800  # 10 * 3 * 50 * 1.2

        cost = TokenEstimator.estimate_cost(1000, 0.002)
        assert cost == 0.002

        cost_default = TokenEstimator.estimate_cost(1000)
        assert cost_default == 0.002
    # endregion test_cost_estimation

    # region test_preview_parse_llm_response [TYPE Function]
    # @PURPOSE: Verify LLM JSON response parsing includes detected_source_language.
    def test_preview_parse_llm_response(self):
        """Parse LLM response should extract translations with detected_source_language per row."""
        response = json.dumps({
            "rows": [
                {"row_id": "0", "translation": "Привет", "detected_source_language": "en"},
                {"row_id": "1", "translation": "Мир"},
            ]
        })
        result = TranslationPreview._parse_llm_response(response, 2)
        assert result == {
            "0": {"translation": "Привет", "detected_source_language": "en"},
            "1": {"translation": "Мир", "detected_source_language": "und"},
        }
    # endregion test_preview_parse_llm_response

    # region test_preview_parse_llm_response_multilang [TYPE Function]
    # @PURPOSE: Verify multi-language LLM response parsing.
    def test_preview_parse_llm_response_multilang(self):
        """Parse multi-language LLM response should extract per-language translations."""
        response = json.dumps({
            "rows": [
                {"row_id": "0", "detected_source_language": "fr", "ru": "текст", "en": "text", "de": "Text"},
                {"row_id": "1", "detected_source_language": "en", "ru": "привет", "en": "hello"},
            ]
        })
        result = TranslationPreview._parse_llm_response(response, 2, target_languages=["ru", "en", "de"])
        assert result == {
            "0": {"detected_source_language": "fr", "ru": "текст", "en": "text", "de": "Text"},
            "1": {"detected_source_language": "en", "ru": "привет", "en": "hello"},
        }
        # Verify per-language extraction
        assert result["0"]["ru"] == "текст"
        assert result["0"]["en"] == "text"
        assert result["1"]["ru"] == "привет"
    # endregion test_preview_parse_llm_response_multilang

    # region test_preview_parse_llm_response_with_code_block [TYPE Function]
    # @PURPOSE: Verify LLM response parsing handles markdown code blocks.
    def test_preview_parse_llm_response_with_code_block(self):
        """Parse LLM response should handle markdown code block wrapping."""
        response = "```json\n{\n    \"rows\": [\n        {\"row_id\": \"0\", \"translation\": \"Test\"}\n    ]\n}\n```"
        result = TranslationPreview._parse_llm_response(response, 1)
        assert result == {"0": {"translation": "Test", "detected_source_language": "und"}}
    # endregion test_preview_parse_llm_response_with_code_block

    # region test_cost_warning [TYPE Function]
    # @PURPOSE: Verify cost warning for large sample sizes.
    def test_cost_warning(self):
        """Cost warning should appear for sample_size > 30."""
        warning = TokenEstimator.check_cost_warning(31, 2, 5000, 0.01)
        assert warning is not None
        assert "Large preview" in warning
        assert "5000" in warning

        no_warning = TokenEstimator.check_cost_warning(10, 2, 1000, 0.002)
        assert no_warning is None
    # endregion test_cost_warning

    # region test_preview_compute_config_hash [TYPE Function]
    # @PURPOSE: Verify config hash computation is deterministic.
    def test_preview_compute_config_hash(self):
        """Config hash should be deterministic and non-empty."""
        job = _make_mock_job()
        hash1 = TranslationPreview._compute_config_hash(job)
        hash2 = TranslationPreview._compute_config_hash(job)
        assert hash1 == hash2
        assert len(hash1) == 16
    # endregion test_preview_compute_config_hash

    # region test_preview_missing_datasource [TYPE Function]
    # @PURPOSE: Verify error when job has no datasource configured.
    def test_preview_missing_datasource(self):
        """Preview should fail if job has no datasource."""
        db = MagicMock()
        config_manager = MagicMock()

        job = _make_mock_job(source_datasource_id=None)
        db.query.return_value.filter.return_value.first.return_value = job

        preview_service = TranslationPreview(db, config_manager, "test-user")
        with pytest.raises(ValueError, match="source datasource"):
            preview_service.preview_rows(job_id="job-123")
    # endregion test_preview_missing_datasource

    # region test_preview_missing_translation_column [TYPE Function]
    # @PURPOSE: Verify error when job has no translation column.
    def test_preview_missing_translation_column(self):
        """Preview should fail if job has no translation column."""
        db = MagicMock()
        config_manager = MagicMock()

        job = _make_mock_job(translation_column=None)
        db.query.return_value.filter.return_value.first.return_value = job

        preview_service = TranslationPreview(db, config_manager, "test-user")
        with pytest.raises(ValueError, match="translation column"):
            preview_service.preview_rows(job_id="job-123")
    # endregion test_preview_missing_translation_column


# region TestAutoDetection [TYPE Class]
# @PURPOSE: Tests for auto-detected source language in preview/execution (T090).
class TestAutoDetection:

    # region test_detect_known_language [TYPE Function]
    # @BRIEF LLM response with known source language is correctly parsed.
    def test_detect_known_language(self):
        """Known source languages (en, fr, de, es) should be detected correctly."""
        test_cases = [
            ("en", "Hello world"),
            ("fr", "Bonjour le monde"),
            ("de", "Hallo Welt"),
            ("es", "Hola mundo"),
        ]
        for lang, text in test_cases:
            response = json.dumps({
                "rows": [{"row_id": "0", "translation": text, "detected_source_language": lang}]
            })
            result = TranslationPreview._parse_llm_response(response, 1)
            assert result["0"]["detected_source_language"] == lang, f"Expected {lang} for '{text}'"
    # endregion test_detect_known_language

    # region test_detect_und_for_mixed_content [TYPE Function]
    # @BRIEF Mixed/ambiguous content returns "und".
    def test_detect_und_for_mixed_content(self):
        """When LLM cannot determine language, 'und' is returned."""
        response = json.dumps({
            "rows": [{"row_id": "0", "translation": "...", "detected_source_language": "und"}]
        })
        result = TranslationPreview._parse_llm_response(response, 1)
        assert result["0"]["detected_source_language"] == "und"
    # endregion test_detect_und_for_mixed_content

    # region test_detect_und_flag_for_review [TYPE Function]
    # @BRIEF "und" rows are flagged for manual review.
    def test_detect_und_flag_for_review(self):
        """Rows with und should be flagged for manual review."""
        response = json.dumps({
            "rows": [
                {"row_id": "0", "translation": "text", "detected_source_language": "en"},
                {"row_id": "1", "translation": "---", "detected_source_language": "und"},
            ]
        })
        result = TranslationPreview._parse_llm_response(response, 2)
        assert result["0"]["detected_source_language"] == "en"
        assert result["1"]["detected_source_language"] == "und"
    # endregion test_detect_und_flag_for_review

    # region test_detect_per_row_independence [TYPE Function]
    # @BRIEF Each row's detected language is independent of other rows.
    def test_detect_per_row_independence(self):
        """Rows in same batch with different languages carry independent detections."""
        response = json.dumps({
            "rows": [
                {"row_id": "0", "detected_source_language": "fr", "ru": "французский", "en": "french"},
                {"row_id": "1", "detected_source_language": "de", "ru": "немецкий", "en": "german"},
            ]
        })
        result = TranslationPreview._parse_llm_response(response, 2, target_languages=["ru", "en"])
        assert result["0"]["detected_source_language"] == "fr"
        assert result["1"]["detected_source_language"] == "de"
        assert result["0"]["ru"] == "французский"
        assert result["1"]["en"] == "german"
    # endregion test_detect_per_row_independence

# endregion TestAutoDetection

# endregion TestTranslationPreview

# endregion TranslationPreviewTests