SIENTIAPDE-1646

Update E2E test report and enhance drift analysis handling - Updated the E2E test report metrics to reflect the latest test results, showing 47 collected tests with all passing. - Removed outdated sections related to failed tests and their causes, streamlining the report. - Implemented a regression fix in the drift analysis to handle empty merged frames, ensuring workflows skip export when no drift metrics are available. - Enhanced the `insert_sample_data` and `insert_sample_prediction` functions to allow customizable timestamps for better test accuracy. - Refactored E2E tests to improve clarity and maintainability, particularly in handling repeat scenarios with distinct timestamps.
2026-05-11 11:49:30 -03:00
parent 10c7e292b9
commit 8d34228d7b
9 changed files with 549 additions and 329 deletions
--- a/e2e/test_predictions_batch_prediction_process.py
+++ b/e2e/test_predictions_batch_prediction_process.py
@@ -3,7 +3,7 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
 """

 from decimal import Decimal
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock

 import numpy as np
 import pandas as pd
@@ -14,32 +14,27 @@ from temporalio.worker import Worker

 from e2e.helpers import (
    assert_continue,
-    load_scenario_input,
+    assert_postgres_unique_violation_in_chain,
+    assert_prediction_row_count,
    assert_repeat,
    assert_stop,
    insert_sample_data,
+    insert_sample_prediction,
+    load_scenario_input,
    make_workflow_id,
    start_and_await_workflow,
 )
 from laborious.activities.activities import Activities
 from laborious.workflows.predictions_batch import PredictionsBatch

+DISTINCT_BATCH_TIMESTAMP = '2024-01-01 13:00:00+00:00'
+HISTORY_TIMESTAMP = '2024-01-01 12:00:00+00:00'
+
+
 def get_base_input_data(model_id):
    return load_scenario_input('prediction_process_base.json', model_id=model_id)


-def insert_sample_prediction(postgres_engine, model_id):
-    with postgres_engine.begin() as conn:
-        conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
-        insert_sql = f"""
-            INSERT INTO sientia_data.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
-            VALUES
-               ({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1)
-        """
-        conn.execute(text(insert_sql))
-    return (model_id, Decimal(10), Decimal(0), 'Good')
-
-
@pytest.fixture
 def bad_data_model(mlflow_repository_stub):
    mlflow_repository_stub.stub_wrapper.transform = MagicMock(
@@ -113,22 +108,52 @@ async def test_scenario_2_1_2_input_gate_triggers_stop(

@pytest.mark.asyncio
@pytest.mark.integration
-async def test_scenario_2_1_3_input_gate_triggers_repeat(
+async def test_scenario_2_1_3_input_gate_repeat_batch_timestamp_equals_history_fails(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
 ):
-    """Input gate REPEAT with existing history."""
+    """
+    REPEAT uses ``last_timestamp`` from the batch payload as the new row's ``timestamp``.
+    When it equals the only historical prediction row, Postgres rejects the duplicate key.
+    """
    client = temporal_test_env.client
    model_id = 213
-    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
-    data = insert_sample_prediction(postgres_engine, model_id)
+    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2], data_timestamp=HISTORY_TIMESTAMP)
+    insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
+    input_data = get_base_input_data(model_id)
+    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
+    with pytest.raises(Exception) as excinfo:
+        await start_and_await_workflow(
+            client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-collision')
+        )
+    assert_postgres_unique_violation_in_chain(excinfo.value)
+    assert_prediction_row_count(postgres_engine, model_id, 1)
+    mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_1_3_input_gate_repeat_distinct_batch_timestamp_inserts_second_row(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    mlflow_repository_stub,
+):
+    """REPEAT succeeds when batch ``last_timestamp`` differs from the historical prediction row."""
+    client = temporal_test_env.client
+    model_id = 2131
+    insert_sample_data(
+        postgres_engine, model_id, ['NULL', 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
+    )
+    data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
    input_data = get_base_input_data(model_id)
    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
    await start_and_await_workflow(
-        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat')
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-ok')
    )
    assert_repeat(postgres_engine, model_id, data)
    mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
@@ -205,7 +230,7 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop(

@pytest.mark.asyncio
@pytest.mark.integration
-async def test_scenario_2_2_3_transform_gate_triggers_repeat(
+async def test_scenario_2_2_3_transform_gate_repeat_batch_timestamp_equals_history_fails(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
@@ -214,12 +239,37 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat(
 ):
    client = temporal_test_env.client
    model_id = 223
-    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
-    data = insert_sample_prediction(postgres_engine, model_id)
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2], data_timestamp=HISTORY_TIMESTAMP)
+    insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
+    with pytest.raises(Exception) as excinfo:
+        await start_and_await_workflow(
+            client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat-collision')
+        )
+    assert_postgres_unique_violation_in_chain(excinfo.value)
+    assert_prediction_row_count(postgres_engine, model_id, 1)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_2_3_transform_gate_repeat_distinct_batch_timestamp_inserts_second_row(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_data_model,
+):
+    client = temporal_test_env.client
+    model_id = 2231
+    insert_sample_data(
+        postgres_engine, model_id, [60.0, 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
+    )
+    data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
    input_data = get_base_input_data(model_id)
    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
    await start_and_await_workflow(
-        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat')
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat-ok')
    )
    assert_repeat(postgres_engine, model_id, data)

@@ -306,7 +356,7 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop(

@pytest.mark.asyncio
@pytest.mark.integration
-async def test_scenario_2_3_3_predict_gate_triggers_repeat(
+async def test_scenario_2_3_3_predict_gate_repeat_batch_timestamp_equals_history_fails(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
@@ -315,13 +365,39 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat(
 ):
    client = temporal_test_env.client
    model_id = 233
-    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-    data = insert_sample_prediction(postgres_engine, model_id)
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2], data_timestamp=HISTORY_TIMESTAMP)
+    insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
+    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
+    with pytest.raises(Exception) as excinfo:
+        await start_and_await_workflow(
+            client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat-collision')
+        )
+    assert_postgres_unique_violation_in_chain(excinfo.value)
+    assert_prediction_row_count(postgres_engine, model_id, 1)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_3_3_predict_gate_repeat_distinct_batch_timestamp_inserts_second_row(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_predict_model,
+):
+    client = temporal_test_env.client
+    model_id = 2331
+    insert_sample_data(
+        postgres_engine, model_id, [23.5, 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
+    )
+    data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
    input_data = get_base_input_data(model_id)
    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
    await start_and_await_workflow(
-        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat')
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat-ok')
    )
    assert_repeat(postgres_engine, model_id, data)