Fix RestException bug (mlflow#1780)

TugceOzberkYener · Aug 26, 2019 · 5a90a6d · 5a90a6d
1 parent 3708f0b
commit 5a90a6d
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/mlflow/exceptions.py b/mlflow/exceptions.py
@@ -56,11 +56,10 @@ def get_http_status_code(self):
 class RestException(MlflowException):
     """Exception thrown on non 200-level responses from the REST API"""
     def __init__(self, json):
-        error_code = json.get('error_code', INTERNAL_ERROR)
+        error_code = json.get('error_code', ErrorCode.Name(INTERNAL_ERROR))
         message = "%s: %s" % (error_code,
                               json['message'] if 'message' in json else "Response: " + str(json))
-
-        super(RestException, self).__init__(message, error_code=error_code)
+        super(RestException, self).__init__(message, error_code=ErrorCode.Value(error_code))
         self.json = json
 
 

diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
@@ -1,6 +1,6 @@
 import json
 
-from mlflow.exceptions import MlflowException
+from mlflow.exceptions import MlflowException, RestException
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, INVALID_STATE, \
     ENDPOINT_NOT_FOUND, INTERNAL_ERROR, RESOURCE_ALREADY_EXISTS, IO_ERROR
 
@@ -33,3 +33,11 @@ def test_get_http_status_code(self):
             == 500
         assert MlflowException('test', error_code=RESOURCE_ALREADY_EXISTS).get_http_status_code() \
             == 400
+
+
+def test_rest_exception():
+    mlflow_exception = MlflowException('test', error_code=RESOURCE_ALREADY_EXISTS)
+    json_exception = mlflow_exception.serialize_as_json()
+    deserialized_rest_exception = RestException(json.loads(json_exception))
+    assert deserialized_rest_exception.error_code == "RESOURCE_ALREADY_EXISTS"
+    assert "test" in deserialized_rest_exception.message
diff --git a/tests/utils/test_exception.py b/tests/utils/test_exception.py
@@ -1,5 +1,6 @@
 import json
 from mlflow.exceptions import ExecutionException, RestException
+from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST, ErrorCode
 
 
 def test_execution_exception_string_repr():
@@ -28,7 +29,8 @@ def test_rest_exception_without_message():
 
 
 def test_rest_exception_error_code_and_no_message():
-    exc = RestException({"error_code": 2, "messages": "something important."})
+    exc = RestException({"error_code": ErrorCode.Name(RESOURCE_DOES_NOT_EXIST),
+                         "messages": "something important."})
     assert "something important." in str(exc)
-    assert "2" in str(exc)
+    assert "RESOURCE_DOES_NOT_EXIST" in str(exc)
     json.loads(exc.serialize_as_json())