QTZPOC-13

QTZPOC-13: Dynamically resolve MLflow pyfunc model code paths to support varied layouts during re-logging.

tests/laborious/utils/repository/test_model_repository.py

@@ -1049,21 +1049,52 @@ async def test_log_model_sklearn(mlflow, mlflow_repository):
     )
 
 
-@patch('laborious.utils.repository.model_repository.path')
+def test_resolve_code_paths_package_layout(tmp_path, mlflow_repository):
+    code_dir = tmp_path / 'code'
+    (code_dir / 'utils').mkdir(parents=True)
+    (code_dir / '__pycache__').mkdir()
+
+    assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [
+        str(code_dir / 'utils')
+    ]
+
+
+def test_resolve_code_paths_single_module_layout(tmp_path, mlflow_repository):
+    code_dir = tmp_path / 'code'
+    code_dir.mkdir()
+    (code_dir / 'loop_valve_classifier_model.py').touch()
+
+    assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [
+        str(code_dir / 'loop_valve_classifier_model.py')
+    ]
+
+
+def test_resolve_code_paths_without_code_dir(tmp_path, mlflow_repository):
+    assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == []
+
+
+def test_resolve_code_paths_without_artifact_path(mlflow_repository):
+    assert mlflow_repository.resolve_code_paths(None, metadata['metadata']) == []
+
+
 @pytest.mark.asyncio
-async def test_log_model_pyfunc(path, mlflow, mlflow_repository):
+async def test_log_model_pyfunc(tmp_path, mlflow, mlflow_repository):
+    code_dir = tmp_path / 'code'
+    code_dir.mkdir()
+    (code_dir / 'loop_valve_classifier_model.py').touch()
+
     model_mock = MagicMock()
-    model_data = {'model': model_mock, 'artifact_path': 'artifact_path'}
+    model_data = {'model': model_mock, 'artifact_path': str(tmp_path)}
     await mlflow_repository.log_model(
         model_data, 'pyfunc', 'prediction_model', metadata['metadata']
     )
 
     mlflow.pyfunc.log_model.assert_not_called()
 
-    path.join.assert_called_once_with('artifact_path', 'code', 'utils')
-
     model_mock.store_model.assert_called_once_with(
-        artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False
+        artifact_path='prediction_model',
+        code_path=[str(code_dir / 'loop_valve_classifier_model.py')],
+        to_disk=False,
     )
 
     mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY)