feat: enhance test scenarios and configuration for regression models
- Updated `pyproject.toml` to include new linting rules for end-to-end tests. - Modified `requirements-dev.txt` to add dependencies for E2E testing with `testcontainers` and `requests`. - Refactored multiple JSON test scenario files to standardize structure, including new fields for `experiment_run_id`, `bucket_name`, and `file_name`. - Improved model training parameters in `train_model_params.py` to use `experiment_name` directly. - Adjusted `data_manager_repository.py` to utilize the updated `experiment_name` for logging. These changes improve the organization and clarity of regression model tests and enhance the overall testing framework.
This commit is contained in:
vitor-aignosi
272
e2e/test_train_model_workflow.py
Normal file
272
e2e/test_train_model_workflow.py
Normal file
@@ -0,0 +1,272 @@
|
||||
"""
|
||||
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)
|
||||
Reference in New Issue
Block a user