gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feature-rc/SIENTIAPDE-2164 - refactor: centralize prediction metadata write in a dedicated activity; REPEAT now writes metadata too

Browse Source
This commit is contained in:
Carla Pinto
2026-10-01 17:46:20 +02:00
parent db32450c68
commit 9c628e70a6
11 changed files with 606 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
tests/laborious/activities/test_activities.py
View File

@@ -6,6 +6,7 @@ from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.model_metrics import ModelMetrics
from laborious.activities.opc import OPC
from laborious.activities.prediction_metadata import PredictionMetadata
from laborious.activities.storage import Storage
@@ -13,6 +14,7 @@ from laborious.activities.storage import Storage
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.PredictionMetadata.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MinioRepository')
@@ -22,6 +24,7 @@ def test___init__(
mock_minio_repository,
mock_api_init,
mock_model_metrics_init,
mock_prediction_metadata_init,
mock_gates_init,
mock_opc_init,
mock_mlflow_init,
@@ -84,6 +87,7 @@ def test___init__(
assert isinstance(activities, Gates)
assert isinstance(activities, ModelMetrics)
assert isinstance(activities, API)
assert isinstance(activities, PredictionMetadata)
mock_storage_init.assert_called_once_with(
ANY,
@@ -122,6 +126,13 @@ def test___init__(
mock_gates_init.assert_called_once_with(
ANY,
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_prediction_metadata_init.assert_called_once_with(
ANY,
mongo_repository=mongo_repository,
logger=logger,
notification_handler=notification_handler,
@@ -161,6 +172,7 @@ def test___init__(
@patch('laborious.activities.activities.MLFlow')
@patch('laborious.activities.activities.OPC')
@patch('laborious.activities.activities.Gates')
@patch('laborious.activities.activities.PredictionMetadata')
@patch('laborious.activities.activities.ModelMetrics')
@patch('laborious.activities.activities.API')
@patch('laborious.activities.activities.MinioRepository')
@@ -168,6 +180,7 @@ def test_shutdown(
_mock_minio_repository,
mock_api_init,
mock_model_metrics_init,
mock_prediction_metadata_init,
mock_gates_init,
mock_opc_init,
mock_mlflow_init,
@@ -230,6 +243,7 @@ def test_shutdown(
mock_gates_init.close.assert_called_once()
mock_model_metrics_init.close.assert_called_once()
mock_api_init.close.assert_called_once()
mock_prediction_metadata_init.close.assert_called_once()
@patch('laborious.activities.activities.SientiaMLflowRepository')
@@ -238,6 +252,7 @@ def test_shutdown(
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.PredictionMetadata.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MinioRepository')
@@ -247,6 +262,7 @@ def test___init___builds_mlflow_repository_when_not_provided(
mock_minio_repository,
_mock_api_init,
_mock_model_metrics_init,
_mock_prediction_metadata_init,
_mock_gates_init,
_mock_opc_init,
_mock_mlflow_init,
@@ -310,6 +326,7 @@ def test___init___builds_mlflow_repository_when_not_provided(
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.PredictionMetadata.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MinioRepository')
@@ -319,6 +336,7 @@ def test___init___builds_mongo_repository_when_not_provided(
mock_minio_repository,
_mock_api_init,
_mock_model_metrics_init,
mock_prediction_metadata_init,
mock_gates_init,
_mock_opc_init,
_mock_mlflow_init,
@@ -373,10 +391,17 @@ def test___init___builds_mongo_repository_when_not_provided(
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
# The same repository built here is the one handed to Gates, not a second instance.
mock_gates_init.assert_called_once_with(
ANY,
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
# The same repository built here is the one handed to PredictionMetadata, not a second
# instance.
mock_prediction_metadata_init.assert_called_once_with(
ANY,
mongo_repository=mock_mongo_repository_cls.return_value,
logger=logger,
notification_handler=notification_handler,

157
tests/laborious/activities/test_gates.py
View File

@@ -35,14 +35,8 @@ def _minio_payload(retrieve_return, status=None):
@fixture
def mongo() -> MagicMock:
return MagicMock()
@fixture
def gates_activity(mongo: MagicMock):
def gates_activity():
gates = Gates(
mongo_repository=mongo,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(spec=MetricsController),
@@ -67,75 +61,6 @@ metadata = {
}
def test_write_prediction_metadata_no_document_skips_mongo(gates_activity, mongo: MagicMock):
# Act
result = gates_activity._write_prediction_metadata(None, metadata)
# Assert
assert result is None
mongo.insert.assert_not_called()
def test_write_prediction_metadata_empty_document_skips_mongo(gates_activity, mongo: MagicMock):
# Act
result = gates_activity._write_prediction_metadata({}, metadata)
# Assert
assert result is None
mongo.insert.assert_not_called()
def test_write_prediction_metadata_not_configured_returns_none():
# Arrange
gates = Gates(
mongo_repository=None,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(spec=MetricsController),
)
gates.error = MagicMock()
# Act
result = gates._write_prediction_metadata({'filter_triggered': 'EMPTY_DATA'}, metadata)
# Assert
assert result is None
gates.error.assert_not_called()
def test_write_prediction_metadata_success_generates_id_and_inserts(
gates_activity, mongo: MagicMock
):
# Arrange
document = {'filter_triggered': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 3}
# Act
result = gates_activity._write_prediction_metadata(document, metadata)
# Assert: a plain 24-hex-char string id (matches predictions.metadata_id's varchar(24)),
# not a BSON ObjectId — `insert(...)` never hands one back, so it's generated client-side.
assert result is not None
assert len(result) == 24
int(result, 16) # raises ValueError if not hex
collection, inserted_document, inserted_metadata = mongo.insert.call_args.args
assert collection == Gates._PREDICTION_METADATA_COLLECTION
assert inserted_document == {**document, '_id': result}
assert inserted_metadata == metadata
def test_write_prediction_metadata_insert_failure_swallowed(gates_activity, mongo: MagicMock):
# Arrange
mongo.insert.side_effect = RuntimeError('connection reset')
# Act
result = gates_activity._write_prediction_metadata({'filter_triggered': 'EMPTY_DATA'}, metadata)
# Assert: the prediction write must never be blocked or failed by this
assert result is None
gates_activity.error.assert_called_once()
def test_input_filter_counters_covers_every_input_filter_function():
# A filter added to input_filter_functions without a matching entry here would degrade to
# counter=0 (see the `.get(..., lambda *_: 0)` fallback in input_gate) instead of crashing —
@@ -143,7 +68,7 @@ def test_input_filter_counters_covers_every_input_filter_function():
assert set(input_filter_counters) == set(input_filter_functions)
def test_input_gate_invalid_filter(gates_activity, mongo: MagicMock):
def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
@@ -160,7 +85,6 @@ def test_input_gate_invalid_filter(gates_activity, mongo: MagicMock):
gates_activity.error.assert_called_once_with(
'Filter INVALID_FILTER not found', metadata['metadata']
)
mongo.insert.assert_not_called()
@patch('laborious.activities.gates.input_filter_functions')
@@ -192,7 +116,7 @@ def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity
)
def test_input_gate_no_filters(gates_activity, mongo: MagicMock):
def test_input_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
@@ -207,13 +131,10 @@ def test_input_gate_no_filters(gates_activity, mongo: MagicMock):
# Assert
assert result == (None, 0, '', None)
gates_activity.debug.assert_called()
mongo.insert.assert_not_called()
def test_input_gate_with_filter(gates_activity, mongo: MagicMock):
# Arrange: STOP never writes to Mongo — see business rule refinement (2026-09-29): STOP
# terminates the workflow before any `predictions` row is written, so a metadata document
# written here would be orphaned.
def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}},
@@ -222,12 +143,15 @@ def test_input_gate_with_filter(gates_activity, mongo: MagicMock):
}
# Act
result = gates_activity.input_gate(input_data)
path_flag, confidence, comment, document = gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, 'Input data with bad quality', None)
# Assert: the gate only computes the document now — it never writes anywhere, and never
# checks path_flag to decide whether to build one (that decision moved to the caller, see
# `PredictionMetadata.write_prediction_metadata` and `prediction_process.py`), so even a STOP
# resolution still gets its document built.
assert (path_flag, confidence, comment) == ('STOP', -1, 'Input data with bad quality')
assert document == {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
gates_activity.debug.assert_called()
mongo.insert.assert_not_called()
def test_input_gate_with_filter_lowercase_keys(gates_activity):
@@ -243,7 +167,12 @@ def test_input_gate_with_filter_lowercase_keys(gates_activity):
result = gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, 'Input data with bad quality', None)
assert result == (
'STOP',
-1,
'Input data with bad quality',
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]},
)
def test_input_gate_with_filter_capitalized_keys(gates_activity):
@@ -259,10 +188,15 @@ def test_input_gate_with_filter_capitalized_keys(gates_activity):
result = gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, 'Input data with bad quality', None)
assert result == (
'STOP',
-1,
'Input data with bad quality',
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]},
)
def test_input_gate_with_filter_not_caught(gates_activity, mongo: MagicMock):
def test_input_gate_with_filter_not_caught(gates_activity):
# Arrange
input_data = {
**metadata,
@@ -277,12 +211,12 @@ def test_input_gate_with_filter_not_caught(gates_activity, mongo: MagicMock):
# Assert
assert result == (None, 0, '', None)
gates_activity.debug.assert_called()
mongo.insert.assert_not_called()
def test_input_gate_repeat_does_not_write_prediction_metadata(gates_activity, mongo: MagicMock):
# Arrange: REPEAT reuses the last existing row instead of writing a new one, so it must not
# write metadata either — same orphaning risk as STOP.
def test_input_gate_repeat_also_builds_document(gates_activity):
# Arrange: REPEAT used to be treated like STOP (no document at all) — it isn't anymore, since
# REPEAT does insert a new `predictions` row (reusing the previous row's values), so it needs
# a document/metadata_id just like CONTINUE (2026-10-01 redesign).
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'POLICY': 'REPEAT', 'CONFIG': {}}},
@@ -294,11 +228,15 @@ def test_input_gate_repeat_does_not_write_prediction_metadata(gates_activity, mo
result = gates_activity.input_gate(input_data)
# Assert
assert result == ('REPEAT', -1, 'Input data with bad quality', None)
mongo.insert.assert_not_called()
assert result == (
'REPEAT',
-1,
'Input data with bad quality',
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'REPEAT', 'counter': 0}]},
)
def test_input_gate_continue_writes_prediction_metadata(gates_activity, mongo: MagicMock):
def test_input_gate_continue_builds_document_with_counter(gates_activity):
# Arrange: two rows with a null in one of the configured variables ('x'); a third row is
# clean, to confirm the counter reflects only the rows that actually matched the condition.
input_data = {
@@ -316,23 +254,16 @@ def test_input_gate_continue_writes_prediction_metadata(gates_activity, mongo: M
}
# Act
path_flag, confidence, comment, metadata_id = gates_activity.input_gate(input_data)
path_flag, confidence, comment, document = gates_activity.input_gate(input_data)
# Assert
assert (path_flag, confidence, comment) == ('CONTINUE', 2, 'Input data with bad quality')
assert metadata_id is not None
collection, inserted_document, inserted_metadata = mongo.insert.call_args.args
assert collection == Gates._PREDICTION_METADATA_COLLECTION
assert inserted_document == {
'filters': [{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}],
'_id': metadata_id,
assert document == {
'filters': [{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}]
}
assert inserted_metadata == metadata['metadata']
def test_input_gate_continue_multiple_filters_same_policy_lists_each(
gates_activity, mongo: MagicMock
):
def test_input_gate_multiple_filters_same_policy_lists_each(gates_activity):
# Arrange: force both input filters to trigger (not reachable with today's real filter
# conditions — EMPTY_DATA and SPECIFIC_VARIABLES_NULL_VALUES are mutually exclusive) to lock
# in the list-building behavior approved by Vitor (2026-09-30): one entry per filter sharing
@@ -355,17 +286,15 @@ def test_input_gate_continue_multiple_filters_same_policy_lists_each(
}
# Act
path_flag, _, _, metadata_id = gates_activity.input_gate(input_data)
path_flag, _, _, document = gates_activity.input_gate(input_data)
# Assert
assert path_flag == 'CONTINUE'
assert metadata_id is not None
_, inserted_document, _ = mongo.insert.call_args.args
assert {f['name'] for f in inserted_document['filters']} == {
assert {f['name'] for f in document['filters']} == {
'SPECIFIC_VARIABLES_NULL_VALUES',
'EMPTY_DATA',
}
assert all(f['policy'] == 'CONTINUE' for f in inserted_document['filters'])
assert all(f['policy'] == 'CONTINUE' for f in document['filters'])
def test_mlflow_response_gate_invalid_filter(gates_activity):

179
tests/laborious/activities/test_prediction_metadata.py Normal file
View File

@@ -0,0 +1,179 @@
from unittest.mock import MagicMock, call, patch
from pytest import fixture
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from laborious.activities.prediction_metadata import PredictionMetadata
@fixture
def mongo() -> MagicMock:
return MagicMock()
@fixture
def prediction_metadata(mongo: MagicMock):
activity = PredictionMetadata(
mongo_repository=mongo,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(spec=MetricsController),
)
activity.error = MagicMock()
return activity
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test_write_prediction_metadata_no_documents_skips_mongo(prediction_metadata, mongo: MagicMock):
# Act
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': []})
# Assert
assert result is None
mongo.insert.assert_not_called()
def test_write_prediction_metadata_missing_documents_key_skips_mongo(
prediction_metadata, mongo: MagicMock
):
# Act
result = prediction_metadata.write_prediction_metadata({**metadata})
# Assert
assert result is None
mongo.insert.assert_not_called()
def test_write_prediction_metadata_documents_with_no_filters_skips_mongo(
prediction_metadata, mongo: MagicMock
):
# Arrange: a falsy document (None) and one with an empty `filters` list — neither contributes
# anything to write.
documents = [None, {'filters': []}]
# Act
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': documents})
# Assert
assert result is None
mongo.insert.assert_not_called()
def test_write_prediction_metadata_not_configured_returns_none():
# Arrange
activity = PredictionMetadata(
mongo_repository=None,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(spec=MetricsController),
)
activity.error = MagicMock()
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
# Act
result = activity.write_prediction_metadata({**metadata, 'documents': [document]})
# Assert
assert result is None
activity.error.assert_not_called()
def test_write_prediction_metadata_success_generates_id_and_inserts(
prediction_metadata, mongo: MagicMock
):
# Arrange
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 3}]}
# Act
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': [document]})
# Assert: a plain 24-hex-char string id (matches predictions.metadata_id's varchar(24)),
# not a BSON ObjectId — `insert(...)` never hands one back, so it's generated client-side.
assert result is not None
assert len(result) == 24
int(result, 16) # raises ValueError if not hex
collection, inserted_document, inserted_metadata = mongo.insert.call_args.args
assert collection == PredictionMetadata._PREDICTION_METADATA_COLLECTION
assert inserted_document == {'filters': document['filters'], '_id': result}
assert inserted_metadata == metadata['metadata']
def test_write_prediction_metadata_combines_multiple_documents(
prediction_metadata, mongo: MagicMock
):
# Arrange: two documents (e.g. from two different sources contributing metadata for the same
# prediction cycle) — their `filters` entries are flattened into one combined document.
documents = [
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'CONTINUE', 'counter': 0}]},
{
'filters': [
{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}
]
},
]
# Act
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': documents})
# Assert
assert result is not None
_, inserted_document, _ = mongo.insert.call_args.args
assert inserted_document['filters'] == [
{'name': 'EMPTY_DATA', 'policy': 'CONTINUE', 'counter': 0},
{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2},
]
def test_write_prediction_metadata_insert_failure_swallowed(prediction_metadata, mongo: MagicMock):
# Arrange
mongo.insert.side_effect = RuntimeError('connection reset')
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
# Act
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': [document]})
# Assert: the prediction write must never be blocked or failed by this
assert result is None
prediction_metadata.error.assert_called_once()
@patch.object(SientiaMonitoring, 'shutdown')
def test_close_disposes_mongo_repository(mock_shutdown):
"""``PredictionMetadata.close`` should close the Mongo client and clear the reference."""
mongo = MagicMock()
activity = PredictionMetadata(
mongo_repository=mongo,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(spec=MetricsController),
)
activity.close()
mongo.close.assert_called_once()
assert activity.mongo_repository is None
# Other instances collected by GC during the test also hit the patched shutdown.
assert mock_shutdown.call_args_list.count(call(activity)) == 1
@patch.object(SientiaMonitoring, 'shutdown')
def test_close_without_mongo_repository(mock_shutdown):
"""When no Mongo repository is configured, ``close`` only shuts down monitoring."""
activity = PredictionMetadata(
mongo_repository=None,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(spec=MetricsController),
)
activity.close()
# Other instances collected by GC during the test also hit the patched shutdown.
assert mock_shutdown.call_args_list.count(call(activity)) == 1

65
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -326,6 +326,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'],
'metadata_id': None,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -389,6 +390,56 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_default_path_flag_passes_metadata_id(
workflow_mock, format_and_export_prediction
):
# Arrange: metadata_id reaches format_default_prediction, which is the row that ends up in
# Postgres via export_data_to_postgres (CONTINUE/REPEAT path).
input_data = {
'metadata': metadata,
'path_flag': 'default',
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'model_name': metadata['metadata']['model_name'],
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'opc_output_config': {'test': 'config'},
'comment': 'test_comment',
'metadata_id': 'a' * 24,
}
workflow_mock.execute_activity_method.side_effect = [
(MagicMock(), MagicMock()),
MagicMock(),
MagicMock(),
]
# Act
await format_and_export_prediction.run(input_data)
# Assert
workflow_mock.execute_local_activity_method.assert_any_call(
Activities.format_default_prediction,
{
**metadata,
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'],
'metadata_id': 'a' * 24,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
@@ -638,6 +689,20 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_any_call(
Activities.format_default_prediction,
{
**metadata,
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'],
'metadata_id': None,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(

74
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -598,6 +598,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {'retention': '30'}
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
# Act
result = await prediction_process.path_flag_handler(
@@ -616,18 +617,22 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
confidence,
last_timestamp,
'',
document,
)
# Assert
# Assert: STOP never reaches a write, even though it now gets a `document` same as CONTINUE
# and REPEAT do — it just never has anywhere to use it.
assert result is True
workflow_mock.execute_local_activity_method.assert_not_called()
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Arrange
# Arrange: REPEAT now writes metadata just like CONTINUE — it does insert a new `predictions`
# row (reusing the previous row's values), so it needs a `metadata_id` too.
data = {'test': 'data'}
path_flag = 'repeat'
confidence = 0.95
@@ -637,6 +642,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {'retention': '30'}
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'REPEAT', 'counter': 0}]}
# Act
result = await prediction_process.path_flag_handler(
@@ -655,11 +661,20 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
confidence,
last_timestamp,
'',
document,
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_called_once_with(
metadata_id = workflow_mock.execute_activity_method.return_value
assert workflow_mock.execute_activity_method.call_count == 2
workflow_mock.execute_activity_method.assert_any_call(
Activities.write_prediction_metadata,
{**metadata, 'documents': [document]},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_activity_method.assert_any_call(
Activities.repeat_last_prediction,
{
**metadata,
@@ -667,6 +682,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
'table_name': table_name,
'model': model,
'last_timestamp': last_timestamp,
'metadata_id': metadata_id,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -674,6 +690,45 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_repeat_no_document_writes_nothing(
workflow_mock, prediction_process
):
# Arrange: no document (e.g. REPEAT triggered by a gate that doesn't build one yet) still
# calls write_prediction_metadata, with an empty `documents` list — the activity itself (not
# this workflow) decides there's nothing to write and returns `None`.
data = {'test': 'data'}
# Act
result = await prediction_process.path_flag_handler(
data,
'repeat',
{
'metadata': metadata,
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'model_id': 'test_model',
'last_timestamp': '2024-01-01',
'model_name': 'test_model_name',
'model_config': {'retention': '30'},
},
0.95,
'2024-01-01',
'',
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_any_call(
Activities.write_prediction_metadata,
{**metadata, 'documents': []},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
@mark.asyncio
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
@@ -688,6 +743,9 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
model_name = 'test_model_name'
model_config = {'retention': '30'}
prediction_store_policy = 'erl:1'
document = {
'filters': [{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}]
}
# Act
result = await prediction_process.path_flag_handler(
@@ -713,11 +771,18 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
confidence,
last_timestamp,
'Prediction Process',
document,
)
# Assert
metadata_id = workflow_mock.execute_activity_method.return_value
assert result is True
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.write_prediction_metadata,
{**metadata, 'documents': [document]},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_child_workflow.assert_called_once_with(
'subworkflow.format_and_export_prediction',
{
@@ -733,6 +798,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'table_name': table_name,
'transform_table_name': 'test_transform_table',
'comment': 'Prediction Process',
'metadata_id': metadata_id,
'opc_output_config': {'test': 'config'},
'pi_web_api_output_config': {
'endpoint': 'https://test-pi-server.com',