Propagate API response error to client for list_artifacts (#5189)

* propagate API response error to client for list_artifacts Signed-off-by: Wendy Hu <wendy.hu@databricks.com> * Add error message to test Signed-off-by: Wendy Hu <wendy.hu@databricks.com> * lint Signed-off-by: Wendy Hu <wendy.hu@databricks.com>
mlflow · Jan 4, 2022 · d6abb30 · d6abb30
1 parent 04e6c82
commit d6abb30
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/mlflow/store/artifact/databricks_models_artifact_repo.py b/mlflow/store/artifact/databricks_models_artifact_repo.py
@@ -71,8 +71,9 @@ def list_artifacts(self, path=None):
             json_body = self._make_json_body(path, page_token)
             response = self._call_endpoint(json_body, REGISTRY_LIST_ARTIFACTS_ENDPOINT)
             try:
+                response.raise_for_status()
                 json_response = json.loads(response.text)
-            except ValueError:
+            except Exception:
                 raise MlflowException(
                     "API request to list files under path `%s` failed with status code %s. "
                     "Response body: %s" % (path, response.status_code, response.text)

diff --git a/tests/store/artifact/test_databricks_models_artifact_repo.py b/tests/store/artifact/test_databricks_models_artifact_repo.py
@@ -142,15 +142,24 @@ def test_init_with_valid_uri_but_no_profile(self, valid_profileless_artifact_uri
                 DatabricksModelsArtifactRepository(valid_profileless_artifact_uri)
 
     def test_list_artifacts(self, databricks_model_artifact_repo):
+        status_code = 200
+
+        def _raise_for_status():
+            if status_code == 404:
+                raise Exception(
+                    "404 Client Error: Not Found for url: https://shard-uri/api/2.0/mlflow/model-versions/list-artifacts?name=model&version=1"
+                )
+
         list_artifact_dir_response_mock = mock.MagicMock()
-        list_artifact_dir_response_mock.status_code = 200
+        list_artifact_dir_response_mock.status_code = status_code
         list_artifact_dir_json_mock = {
             "files": [
                 {"path": "MLmodel", "is_dir": False, "file_size": 294},
                 {"path": "data", "is_dir": True, "file_size": None},
             ]
         }
         list_artifact_dir_response_mock.text = json.dumps(list_artifact_dir_json_mock)
+        list_artifact_dir_response_mock.raise_for_status.side_effect = _raise_for_status
         with mock.patch(
             DATABRICKS_MODEL_ARTIFACT_REPOSITORY + "._call_endpoint"
         ) as call_endpoint_mock:
@@ -166,6 +175,24 @@ def test_list_artifacts(self, databricks_model_artifact_repo):
             assert artifacts[1].file_size is None
             call_endpoint_mock.assert_called_once_with(ANY, REGISTRY_LIST_ARTIFACTS_ENDPOINT)
 
+        # errors from API are propagated through to cli response
+        list_artifact_dir_bad_response_mock = mock.MagicMock()
+        status_code = 404
+        list_artifact_dir_bad_response_mock.status_code = status_code
+        list_artifact_dir_bad_response_mock.text = "An error occurred"
+        list_artifact_dir_bad_response_mock.raise_for_status.side_effect = _raise_for_status
+        with mock.patch(
+            DATABRICKS_MODEL_ARTIFACT_REPOSITORY + "._call_endpoint"
+        ) as call_endpoint_mock:
+            call_endpoint_mock.return_value = list_artifact_dir_bad_response_mock
+            with pytest.raises(
+                MlflowException,
+                match=r"API request to list files under path `` failed with status code 404. "
+                "Response body: An error occurred",
+            ):
+                databricks_model_artifact_repo.list_artifacts("")
+            call_endpoint_mock.assert_called_once_with(ANY, REGISTRY_LIST_ARTIFACTS_ENDPOINT)
+
     def test_list_artifacts_for_single_file(self, databricks_model_artifact_repo):
         list_artifact_file_response_mock = mock.MagicMock()
         list_artifact_file_response_mock.status_code = 200