# #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 = 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