From 36f13883016b5c3118284e3bbc8ff460be7e9cbf Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 28 Aug 2025 16:32:06 -0300 Subject: [PATCH] SIENTIAPDE-1182 SIENTIAPDE-1205: Enhance predictions_batch function and update test cases - Added 'predictions_storage_policy' to the predictions_batch function for improved configuration handling. - Updated test cases to include 'predictions_storage_policy' in expected results, ensuring consistency in testing. --- orchestrator/utils/orchestrator_functions.py | 3 ++- tests/orchestrator/utils/test_orchestrator_functions.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/orchestrator/utils/orchestrator_functions.py b/orchestrator/utils/orchestrator_functions.py index b49a771..90bfdfb 100644 --- a/orchestrator/utils/orchestrator_functions.py +++ b/orchestrator/utils/orchestrator_functions.py @@ -211,7 +211,8 @@ def predictions_batch(config: dict[str, Any]): "config": {} } }, config.get('mlflow_predict_filters', [])), - "path_priority": path_priority + "path_priority": path_priority, + "predictions_storage_policy": config.get('predictions_storage_policy', 'lts:1') } diff --git a/tests/orchestrator/utils/test_orchestrator_functions.py b/tests/orchestrator/utils/test_orchestrator_functions.py index 09fae10..cd59478 100644 --- a/tests/orchestrator/utils/test_orchestrator_functions.py +++ b/tests/orchestrator/utils/test_orchestrator_functions.py @@ -197,7 +197,8 @@ def test_predictions_batch(mock_process_path_priority, } ], "path_priority": ["STOP", "CONTINUE", "REPEAT"], - "datetime_columns": ["timestamp"] + "datetime_columns": ["timestamp"], + "predictions_storage_policy": "erl:1" } result = predictions_batch(config) @@ -272,7 +273,8 @@ def test_predictions_batch(mock_process_path_priority, } }, "path_priority": ["STOP", "CONTINUE", "REPEAT"], - "datetime_columns": ["timestamp"] + "datetime_columns": ["timestamp"], + "predictions_storage_policy": "erl:1" } assert result == expected