QTZPOC-20

QTZPOC-20: Prioritize code/ directory download and sys.path injection before model loading, enabling MLflow flavors to stream models directly from URIs and simplifying joblib artifact handling.

e2e/conftest.py

@@ -546,14 +546,7 @@ def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts):
     mock_mlflow.pyfunc = MagicMock()
     mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
 
-    # The sklearn and pyfunc paths resolve the URI to a local directory before loading, so that the
-    # artifact's `code/` can go on `sys.path` whatever the `MLmodel` manifest says. Echoing the URI
-    # keeps the loader stubs above — which branch on the URI text — working unchanged, and points
-    # the code-directory lookup at a path that does not exist.
     mock_mlflow.artifacts = MagicMock()
-    mock_mlflow.artifacts.download_artifacts = MagicMock(
-        side_effect=lambda artifact_uri: str(artifact_uri)
-    )
 
     # Mock other mlflow methods that might be called
     mock_mlflow.set_tracking_uri = MagicMock()
@@ -572,9 +565,23 @@ def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts):
         ),
     )
 
-    # Only the joblib flavor calls download_artifacts directly on the client, one artifact
-    # per model type — a single return value would serve `data_model` from `prediction_model`.
+    # `download_model` asks whether the run's artifact carries a `code` directory before it loads
+    # anything, and downloads only that directory when it does. Answering "no" here keeps the
+    # download counts of these tests about the model itself; a test that exercises the `code` step
+    # overrides `list_artifacts` to answer "yes".
+    def mock_list_artifacts(_run_id, artifact_path=None):
+        if artifact_path:
+            return []
+        return [MagicMock(path=artifact_name) for artifact_name in joblib_flavor_artifacts]
+
+    mock_mlflow_client.list_artifacts = MagicMock(side_effect=mock_list_artifacts)
+
+    # The joblib flavor downloads a whole model artifact, one per model type — a single return
+    # value would serve `data_model` from `prediction_model`. The `code` step asks for the nested
+    # `<artifact>/code` instead, and gets that directory back, as MLflow would return it.
     def mock_download_artifacts(_run_id, artifact_name):
+        if artifact_name.endswith('/code'):
+            return joblib_flavor_artifacts[artifact_name.removesuffix('/code')]['code_dir']
         return joblib_flavor_artifacts[artifact_name]['path']
 
     mock_mlflow_client.download_artifacts = MagicMock(side_effect=mock_download_artifacts)