""" End-to-end tests for TrainModel workflow – main workflow scenarios. Covers: 1.1.x – Happy-path training (various scenarios from docs/test-scenarios/) 1.2.x – Error paths (MinIO failure, missing DB row) """ import pytest import pytest_asyncio from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker from e2e.helpers import ( assert_experiment_error, assert_experiment_run_name_set, assert_experiment_status, insert_experiment_run, load_scenario, make_workflow_id, start_and_await_workflow, ) from model_manager.workflows.train_model import TrainModel # --------------------------------------------------------------------------- # 1.1 – Happy paths # --------------------------------------------------------------------------- @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_1_linear_regression_basic( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.1 – Linear Regression Basic (cenário 01). Validates the complete training pipeline end-to-end: load_model_metadata → validate_train_params → train_model → update_experiment_run (TRAINING_SUCCESS). """ scenario = load_scenario('01-linear-regression-basic.json') experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-1'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') assert_experiment_run_name_set(postgres_engine, experiment_run_id) @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_2_polynomial_regression_degree2_with_scaler( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.2 – Polynomial Regression Degree 2 with Standard Scaler (cenário 03).""" scenario = load_scenario('03-polynomial-regression-degree2.json') experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-2'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') assert_experiment_run_name_set(postgres_engine, experiment_run_id) @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_3_linear_regression_with_lags( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.3 – Linear Regression with lag_train/lag_val per variable (cenário 05).""" scenario = load_scenario('05-linear-regression-with-lags.json') experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-3'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') assert_experiment_run_name_set(postgres_engine, experiment_run_id) @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_4_linear_regression_nan_interpolation( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.4 – nan_treatment='linear interpolation' (cenário 06).""" scenario = load_scenario('06-linear-regression-nan-interpolation.json') experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-4'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_5_linear_regression_with_limits( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.5 – support_filters with min/max limits per variable (cenário 08).""" scenario = load_scenario('08-linear-regression-with-limits.json') experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-5'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_6_polynomial_degree2_scaler_and_lags( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.6 – Polynomial degree 2, Standard Scaler and lags (cenário 09).""" scenario = load_scenario('09-polynomial-degree2-with-scaler-and-lags.json') experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-6'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') assert_experiment_run_name_set(postgres_engine, experiment_run_id) @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_7_static_window_removal( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.7 – rem_static_win=true with window and static_threshold (cenário 11).""" scenario = load_scenario('11-linear-regression-static-threshold-custom.json') experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-7'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_1_8_polynomial_with_support_filters( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.1.8 – Polynomial degree 4, Standard Scaler, upper/lower support filters (cenário 14).""" scenario = load_scenario('14-angular-test-polynomial-support-filters.json') # Override date range to match rows in our test CSV scenario['data_model_kwargs']['start_date'] = '2025-06-02 00:00:00' scenario['data_model_kwargs']['end_date'] = '2025-06-06 23:59:59' experiment_run_id = scenario['experiment_run_id'] insert_experiment_run(postgres_engine, experiment_run_id) await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-1-8'), ) assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS') # --------------------------------------------------------------------------- # 1.2 – Error paths # --------------------------------------------------------------------------- @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_2_1_minio_file_not_found( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.2.1 – Training file does not exist in MinIO → TRAINING_ERROR.""" scenario = load_scenario('01-linear-regression-basic.json') scenario = {**scenario, 'experiment_run_id': 2001, 'file_name': 'does_not_exist.csv'} experiment_run_id = 2001 insert_experiment_run(postgres_engine, experiment_run_id) with pytest.raises(Exception): await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-2-1'), ) assert_experiment_error( postgres_engine, experiment_run_id, expected_status='TRAINING_ERROR', error_substr='does_not_exist', ) @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_1_2_2_experiment_run_id_not_in_db( temporal_test_env: WorkflowEnvironment, temporal_worker: Worker, postgres_engine, ): """Scenario 1.2.2 – experiment_run_id row absent → update_experiment_run raises.""" scenario = load_scenario('01-linear-regression-basic.json') scenario = {**scenario, 'experiment_run_id': 9999} # Intentionally NOT inserting the row with pytest.raises(Exception): await start_and_await_workflow( temporal_test_env.client, TrainModel.run, scenario, make_workflow_id('test-s1-2-2'), ) from e2e.helpers import assert_no_experiment_row assert_no_experiment_row(postgres_engine, 9999)