- Made `date_column` a required field in `TrainModelParams`, ensuring it must be present in the input data. - Updated related documentation in `input-sample.md`, `README.md`, and various test scenarios to reflect the change in requirement. - Adjusted the handling of `date_format` to default to `yyyy-MM-dd HH:mm:ss` if omitted, enhancing usability. - Refined test scenarios to include new examples and ensure compliance with the updated parameter structure. These changes improve the robustness of the model training workflow and clarify the expectations for input data.
494 lines
16 KiB
Python
494 lines
16 KiB
Python
"""
|
||
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
|
||
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,
|
||
assert_no_experiment_row,
|
||
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_linear_regression_with_scaler(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.2 – Linear regression with Standard Scaler (cenário 02)."""
|
||
scenario = load_scenario('02-linear-regression-with-scaler.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_polynomial_regression_degree2_with_scaler(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.3 – 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-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_polynomial_regression_degree3_with_scaler(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.4 – Polynomial regression degree 3 with Standard Scaler (cenário 04)."""
|
||
scenario = load_scenario('04-polynomial-regression-degree3.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')
|
||
assert_experiment_run_name_set(postgres_engine, experiment_run_id)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.integration
|
||
async def test_scenario_1_1_5_linear_regression_with_lags(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.5 – 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-5'),
|
||
)
|
||
|
||
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_6_linear_regression_nan_interpolation(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.6 – 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-6'),
|
||
)
|
||
|
||
assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.integration
|
||
async def test_scenario_1_1_7_linear_regression_static_window_removal(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.7 – rem_static_win=true with default static_threshold (cenário 07)."""
|
||
scenario = load_scenario('07-linear-regression-static-window-removal.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_linear_regression_with_limits(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.8 – 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-8'),
|
||
)
|
||
|
||
assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.integration
|
||
async def test_scenario_1_1_9_polynomial_degree2_scaler_and_lags(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.9 – 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-9'),
|
||
)
|
||
|
||
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_10_linear_regression_with_ar_opt_params(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.10 – Linear regression with include_ar in opt_params (cenário 10)."""
|
||
scenario = load_scenario('10-linear-regression-with-ar.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-10'),
|
||
)
|
||
|
||
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_11_linear_regression_static_threshold_custom(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.11 – rem_static_win=true with custom 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-11'),
|
||
)
|
||
|
||
assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.integration
|
||
async def test_scenario_1_1_12_alternate_date_format_dd_mm_yyyy(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.12 – DATA column and dd/MM/yyyy format CSV (cenário 12)."""
|
||
scenario = load_scenario('12-angular-test-date-format.json')
|
||
experiment_run_id = scenario['experiment_run_id']
|
||
insert_experiment_run(
|
||
postgres_engine,
|
||
experiment_run_id,
|
||
file_name='training_data_dd_mm_yyyy.csv',
|
||
)
|
||
|
||
await start_and_await_workflow(
|
||
temporal_test_env.client,
|
||
TrainModel.run,
|
||
scenario,
|
||
make_workflow_id('test-s1-1-12'),
|
||
)
|
||
|
||
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_13_alternate_csv_narrow_date_window(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.13 – Same alternate CSV as 12 with date window (cenário 13)."""
|
||
scenario = load_scenario('13-angular-test-double-date-column.json')
|
||
experiment_run_id = scenario['experiment_run_id']
|
||
insert_experiment_run(
|
||
postgres_engine,
|
||
experiment_run_id,
|
||
file_name='training_data_dd_mm_yyyy.csv',
|
||
)
|
||
|
||
await start_and_await_workflow(
|
||
temporal_test_env.client,
|
||
TrainModel.run,
|
||
scenario,
|
||
make_workflow_id('test-s1-1-13'),
|
||
)
|
||
|
||
assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.integration
|
||
async def test_scenario_1_1_14_polynomial_with_support_filters(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.14 – Polynomial degree 4, Standard Scaler, line support filters (cenário 14)."""
|
||
scenario = load_scenario('14-angular-test-polynomial-support-filters.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-14'),
|
||
)
|
||
|
||
assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.integration
|
||
async def test_scenario_1_1_15_linear_regression_custom_target_column_name(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.15 – Target column is not named ``target``; report path uses target_variable."""
|
||
scenario = load_scenario('15-linear-regression-custom-target-column.json')
|
||
experiment_run_id = scenario['experiment_run_id']
|
||
insert_experiment_run(
|
||
postgres_engine,
|
||
experiment_run_id,
|
||
file_name='training_data_custom_target.csv',
|
||
)
|
||
|
||
await start_and_await_workflow(
|
||
temporal_test_env.client,
|
||
TrainModel.run,
|
||
scenario,
|
||
make_workflow_id('test-s1-1-15'),
|
||
)
|
||
|
||
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_16_naive_timestamp_header_column(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.16 – CSV uses ``Timestamp`` header; snake_case date_column/date_format."""
|
||
scenario = load_scenario('16-linear-regression-naive-timestamp-header.json')
|
||
experiment_run_id = scenario['experiment_run_id']
|
||
insert_experiment_run(
|
||
postgres_engine,
|
||
experiment_run_id,
|
||
file_name='training_data_timestamp_naive.csv',
|
||
)
|
||
|
||
await start_and_await_workflow(
|
||
temporal_test_env.client,
|
||
TrainModel.run,
|
||
scenario,
|
||
make_workflow_id('test-s1-1-16'),
|
||
)
|
||
|
||
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_17_linear_regression_blank_timestamp_row_dropped(
|
||
temporal_test_env: WorkflowEnvironment,
|
||
temporal_worker: Worker,
|
||
postgres_engine,
|
||
):
|
||
"""Scenario 1.1.17 – CSV has one empty timestamp cell; row is dropped and training succeeds."""
|
||
scenario = load_scenario('17-linear-regression-blank-timestamp-row.json')
|
||
experiment_run_id = scenario['experiment_run_id']
|
||
insert_experiment_run(
|
||
postgres_engine,
|
||
experiment_run_id,
|
||
file_name='training_data_blank_timestamp_row.csv',
|
||
)
|
||
|
||
await start_and_await_workflow(
|
||
temporal_test_env.client,
|
||
TrainModel.run,
|
||
scenario,
|
||
make_workflow_id('test-s1-1-17'),
|
||
)
|
||
|
||
assert_experiment_status(postgres_engine, experiment_run_id, 'TRAINING_SUCCESS')
|
||
assert_experiment_run_name_set(postgres_engine, experiment_run_id)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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'),
|
||
)
|
||
|
||
assert_no_experiment_row(postgres_engine, 9999)
|