From 66193cea15343a53d606e62e9c9378463c82ed80 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 14 Nov 2025 08:47:39 -0300 Subject: [PATCH] SIENTIAPDE-1273 Refactor Drift class to improve target name handling - Extracted target name from model configuration in the Drift class for better clarity and maintainability. - Updated test cases to reflect changes in input data structure, ensuring consistency in target name retrieval. --- laborious/workflows/drift.py | 5 ++++- tests/laborious/workflows/test_drift.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 5677b37..d7c213c 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -32,6 +32,9 @@ class Drift: } } + model_config = input_data['model_config'] + target_name = model_config['target'] + gathering_query = f""" SELECT * FROM {input_data['schema']}.{input_data['source_table_name']} @@ -76,7 +79,7 @@ class Drift: 'reference_data': reference_data, 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], - 'target_name': input_data['target_name'], + 'target_name': target_name, 'drift_metrics': input_data['drift_metrics'], 'chunk_period': input_data.get('chunk_period', 'min'), }, diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py index 09411aa..82c9eb3 100644 --- a/tests/laborious/workflows/test_drift.py +++ b/tests/laborious/workflows/test_drift.py @@ -34,7 +34,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): 'source_table_name': 'test_source_table', 'target_table_name': 'test_target_table', 'interval': 60, - 'target_name': 'test_target', + 'model_config': {'target': 'test_target'}, 'drift_metrics': ['psi', 'ks'], 'chunk_period': 'hour', }