75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
"""
|
|
Direct E2E execution of child workflows (smaller surface than PredictionsBatch).
|
|
"""
|
|
|
|
from decimal import Decimal
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.helpers import make_workflow_id, start_and_await_workflow
|
|
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_format_and_export_prediction_default_path_e2e(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Run FormatAndExportPrediction with path_flag set (format_default_prediction path).
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 401
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': model_id,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'subworkflow.format_and_export_prediction',
|
|
}
|
|
}
|
|
input_data = {
|
|
'metadata': metadata,
|
|
'path_flag': 'CONTINUE',
|
|
'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'},
|
|
'prediction_confidence': 2,
|
|
'timestamp': '2024-01-01 12:00:00+00:00',
|
|
'model_id': model_id,
|
|
'model_name': 'test_model',
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
'comment': 'e2e child workflow default path',
|
|
'opc_output_config': {},
|
|
'pi_web_api_output_config': {},
|
|
'prediction_store_policy': 'lts:1',
|
|
}
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
FormatAndExportPrediction.run,
|
|
input_data,
|
|
make_workflow_id('e2e-format-export-child'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
row = conn.execute(
|
|
text(
|
|
f'SELECT prediction, prediction_confidence, prediction_status, comments '
|
|
f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
|
|
)
|
|
).fetchone()
|
|
assert row is not None
|
|
assert row[0] == 0
|
|
assert row[1] == Decimal(2)
|
|
assert row[2] == 'Bad'
|
|
assert row[3] == 'e2e child workflow default path'
|