SIENTIAPDE-1646

Enhance MLFlow reference data handling and testing

- Updated the artifact handling in MLFlow to prioritize 'retrain_input.csv' over 'train_data.csv' when resolving reference data.
- Introduced new methods for resolving artifact names and locating downloaded CSV files.
- Modified the `get_reference_data` method to improve artifact resolution and error handling.
- Expanded unit tests to cover scenarios for missing artifacts and preference logic between retrain and train data CSVs.

tests/laborious/activities/test_mlflow.py

@@ -652,6 +652,12 @@ def test_update_production_model_error(mlflow):
         raise AssertionError('Expected exception')
 
 
+def _artifact_file_info(path: str) -> MagicMock:
+    file_info = MagicMock()
+    file_info.path = path
+    return file_info
+
+
 @patch('laborious.activities.mlflow.to_datetime')
 def test_get_reference_data_success(mock_to_datetime, mlflow):
     input_data = {
@@ -661,6 +667,9 @@ def test_get_reference_data_success(mock_to_datetime, mlflow):
 
     mv = MagicMock(run_id='run1')
     mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('retrain_input.csv'),
+    ]
 
     mock_reference_data = MagicMock()
     mock_reference_data.__getitem__.return_value = MagicMock()
@@ -672,10 +681,19 @@ def test_get_reference_data_success(mock_to_datetime, mlflow):
     with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
         with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
             with patch('laborious.activities.mlflow.rmtree'):
-                with patch('laborious.activities.mlflow.Path') as mp:
-                    mp.return_value.rglob.return_value = [MagicMock()]
+                with patch.object(
+                    mlflow,
+                    '_find_downloaded_csv',
+                    return_value=MagicMock(),
+                ):
                     result = mlflow.get_reference_data(input_data)
 
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='retrain_input.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
     mock_reference_data.to_dict.assert_called_once_with(orient='records')
     assert result == mock_reference_data.to_dict.return_value
 
@@ -694,6 +712,97 @@ def test_get_reference_data_not_found(mlflow):
     assert result is None
 
 
+def test_get_reference_data_only_train_data_csv(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('train_data.csv'),
+    ]
+
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
+
+    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
+        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
+            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
+                with patch('laborious.activities.mlflow.rmtree'):
+                    with patch.object(
+                        mlflow,
+                        '_find_downloaded_csv',
+                        return_value=MagicMock(),
+                    ):
+                        result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='train_data.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
+    assert result == mock_reference_data.to_dict.return_value
+
+
+def test_get_reference_data_prefers_retrain_input_when_both_listed(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('train_data.csv'),
+        _artifact_file_info('retrain_input.csv'),
+    ]
+
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
+
+    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
+        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
+            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
+                with patch('laborious.activities.mlflow.rmtree'):
+                    with patch.object(
+                        mlflow,
+                        '_find_downloaded_csv',
+                        return_value=MagicMock(),
+                    ):
+                        result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='retrain_input.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
+    assert result == mock_reference_data.to_dict.return_value
+
+
+def test_get_reference_data_no_candidate_artifacts_returns_none(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('other_artifact.csv'),
+    ]
+
+    result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_not_called()
+    mlflow.warning.assert_called()
+    assert result is None
+
+
 def test_get_reference_data_missing_csv_file_returns_none(mlflow):
     input_data = {
         **metadata,
@@ -701,13 +810,17 @@ def test_get_reference_data_missing_csv_file_returns_none(mlflow):
     }
     mv = MagicMock(run_id='run1')
     mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('retrain_input.csv'),
+    ]
 
     with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='tmp'):
         with patch('laborious.activities.mlflow.rmtree'):
-            with patch('laborious.activities.mlflow.Path') as mp:
-                mp.return_value.rglob.return_value = []
+            with patch.object(mlflow, '_find_downloaded_csv', return_value=None):
                 result = mlflow.get_reference_data(input_data)
 
+    mlflow.mlflow_repository.download_artifacts.assert_called_once()
+    mlflow.warning.assert_called()
     assert result is None
 
 
@@ -719,6 +832,9 @@ def test_get_reference_data_exception(mlflow):
 
     mv = MagicMock(run_id='run1')
     mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('retrain_input.csv'),
+    ]
     mlflow.mlflow_repository.download_artifacts.side_effect = Exception('dl fail')
 
     result = mlflow.get_reference_data(input_data)