wehub-resource-sync
213 lines
7.7 KiB
Python
213 lines
7.7 KiB
Python
import json
|
|
import pickle
|
|
|
|
import pytest
|
|
|
|
from mlflow.exceptions import MlflowException, RestException
|
|
from mlflow.protos.databricks_pb2 import (
|
|
ENDPOINT_NOT_FOUND,
|
|
INTERNAL_ERROR,
|
|
INVALID_PARAMETER_VALUE,
|
|
INVALID_STATE,
|
|
IO_ERROR,
|
|
RESOURCE_ALREADY_EXISTS,
|
|
)
|
|
|
|
|
|
def test_error_code_constructor():
|
|
assert (
|
|
MlflowException("test", error_code=INVALID_PARAMETER_VALUE).error_code
|
|
== "INVALID_PARAMETER_VALUE"
|
|
)
|
|
|
|
|
|
def test_default_error_code():
|
|
assert MlflowException("test").error_code == "INTERNAL_ERROR"
|
|
|
|
|
|
def test_serialize_to_json():
|
|
mlflow_exception = MlflowException("test")
|
|
deserialized = json.loads(mlflow_exception.serialize_as_json())
|
|
assert deserialized["message"] == "test"
|
|
assert deserialized["error_code"] == "INTERNAL_ERROR"
|
|
|
|
|
|
def test_get_http_status_code():
|
|
assert MlflowException("test default").get_http_status_code() == 500
|
|
assert MlflowException("code not in map", error_code=IO_ERROR).get_http_status_code() == 500
|
|
assert MlflowException("test", error_code=INVALID_STATE).get_http_status_code() == 500
|
|
assert MlflowException("test", error_code=ENDPOINT_NOT_FOUND).get_http_status_code() == 404
|
|
assert MlflowException("test", error_code=INVALID_PARAMETER_VALUE).get_http_status_code() == 400
|
|
assert MlflowException("test", error_code=INTERNAL_ERROR).get_http_status_code() == 500
|
|
assert MlflowException("test", error_code=RESOURCE_ALREADY_EXISTS).get_http_status_code() == 400
|
|
|
|
|
|
def test_invalid_parameter_value():
|
|
mlflow_exception = MlflowException.invalid_parameter_value("test")
|
|
assert mlflow_exception.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
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
|
|
|
|
|
|
def test_rest_exception_with_unrecognized_error_code():
|
|
# Test that we can create a RestException with a convertible error code.
|
|
exception = RestException({"error_code": "403", "messages": "something important."})
|
|
assert "something important." in str(exception)
|
|
assert exception.error_code == "PERMISSION_DENIED"
|
|
json.loads(exception.serialize_as_json())
|
|
|
|
# Test that we can create a RestException with an unrecognized error code.
|
|
exception = RestException({"error_code": "weird error", "messages": "something important."})
|
|
assert "something important." in str(exception)
|
|
json.loads(exception.serialize_as_json())
|
|
|
|
|
|
def test_rest_exception_pickleable():
|
|
e1 = RestException({"error_code": "INTERNAL_ERROR", "message": "abc"})
|
|
e2 = pickle.loads(pickle.dumps(e1))
|
|
|
|
assert e1.error_code == e2.error_code
|
|
assert e1.message == e2.message
|
|
|
|
|
|
def test_rest_exception_with_null_error_code():
|
|
exception = RestException({"error_code": None, "message": "test message"})
|
|
assert exception.error_code == "INTERNAL_ERROR"
|
|
assert "test message" in str(exception)
|
|
|
|
|
|
def test_rest_exception_with_missing_error_code():
|
|
exception = RestException({"message": "test message"})
|
|
assert exception.error_code == "INTERNAL_ERROR"
|
|
assert "test message" in str(exception)
|
|
|
|
|
|
# --- sqlstate / error_class auto-derive tests ---
|
|
|
|
|
|
def test_sqlstate_auto_derived_from_error_code():
|
|
exc = MlflowException("test")
|
|
assert exc.sqlstate == "XXM00"
|
|
assert exc.error_class == "CLIENT_INTERNAL_ERROR"
|
|
|
|
exc = MlflowException("test", error_code=INVALID_PARAMETER_VALUE)
|
|
assert exc.sqlstate == "KAM00"
|
|
assert exc.error_class == "INVALID_PARAMETER_VALUE"
|
|
|
|
exc = MlflowException("test", error_code=INTERNAL_ERROR)
|
|
assert exc.sqlstate == "XXM00"
|
|
assert exc.error_class == "CLIENT_INTERNAL_ERROR"
|
|
|
|
|
|
def test_sqlstate_explicit_overrides_auto_derive():
|
|
exc = MlflowException(
|
|
"test",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
sqlstate="KAM01",
|
|
error_class="SCHEMA_ENFORCEMENT_FAILED",
|
|
)
|
|
assert exc.sqlstate == "KAM01"
|
|
assert exc.error_class == "SCHEMA_ENFORCEMENT_FAILED"
|
|
|
|
|
|
def test_sqlstate_serialize_as_json_includes_auto_derived():
|
|
exc = MlflowException("test", error_code=INVALID_PARAMETER_VALUE)
|
|
deserialized = json.loads(exc.serialize_as_json())
|
|
assert deserialized["sqlstate"] == "KAM00"
|
|
assert deserialized["error_class"] == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_sqlstate_serialize_as_json_includes_explicit():
|
|
exc = MlflowException("test", sqlstate="KAM01", error_class="SCHEMA_ENFORCEMENT_FAILED")
|
|
deserialized = json.loads(exc.serialize_as_json())
|
|
assert deserialized["sqlstate"] == "KAM01"
|
|
assert deserialized["error_class"] == "SCHEMA_ENFORCEMENT_FAILED"
|
|
|
|
|
|
def test_sqlstate_none_for_unknown_error_code():
|
|
exc = MlflowException("test", error_code=IO_ERROR)
|
|
assert exc.sqlstate is None
|
|
assert exc.error_class is None
|
|
deserialized = json.loads(exc.serialize_as_json())
|
|
assert "sqlstate" not in deserialized
|
|
assert "error_class" not in deserialized
|
|
|
|
|
|
def test_invalid_parameter_value_auto_derives_sqlstate():
|
|
exc = MlflowException.invalid_parameter_value("bad input")
|
|
assert exc.error_code == "INVALID_PARAMETER_VALUE"
|
|
assert exc.sqlstate == "KAM00"
|
|
assert exc.error_class == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_invalid_parameter_value_with_explicit_override():
|
|
exc = MlflowException.invalid_parameter_value(
|
|
"bad input", sqlstate="KAM02", error_class="PREDICTION_FUNCTION_FAILED"
|
|
)
|
|
assert exc.error_code == "INVALID_PARAMETER_VALUE"
|
|
assert exc.sqlstate == "KAM02"
|
|
assert exc.error_class == "PREDICTION_FUNCTION_FAILED"
|
|
|
|
|
|
# --- RestException CP mapping tests ---
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("error_code", "expected_sqlstate", "expected_error_class"),
|
|
[
|
|
("PERMISSION_DENIED", "KAMC1", "CP_PERMISSION_DENIED"),
|
|
("RESOURCE_DOES_NOT_EXIST", "KAMC2", "CP_RESOURCE_NOT_FOUND"),
|
|
("REQUEST_LIMIT_EXCEEDED", "KAMC3", "CP_REQUEST_RATE_LIMITED"),
|
|
("INVALID_PARAMETER_VALUE", "KAMC4", "CP_INVALID_PARAMETER_VALUE"),
|
|
("INTERNAL_ERROR", "XXMC0", "CP_INTERNAL_ERROR"),
|
|
("TEMPORARILY_UNAVAILABLE", "XXMC1", "CP_TEMPORARILY_UNAVAILABLE"),
|
|
("INVALID_STATE", "XXMC2", "CP_INVALID_STATE"),
|
|
],
|
|
)
|
|
def test_rest_exception_cp_sqlstate_mapping(error_code, expected_sqlstate, expected_error_class):
|
|
exc = RestException({"error_code": error_code, "message": "test"})
|
|
assert exc.sqlstate == expected_sqlstate
|
|
assert exc.error_class == expected_error_class
|
|
|
|
|
|
def test_rest_exception_preserves_sqlstate_from_json():
|
|
exc = RestException({
|
|
"error_code": "PERMISSION_DENIED",
|
|
"message": "no access",
|
|
"sqlstate": "CUSTOM",
|
|
"error_class": "CUSTOM_CLASS",
|
|
})
|
|
assert exc.sqlstate == "CUSTOM"
|
|
assert exc.error_class == "CUSTOM_CLASS"
|
|
|
|
|
|
def test_rest_exception_ignores_null_sqlstate_from_json():
|
|
exc = RestException({
|
|
"error_code": "PERMISSION_DENIED",
|
|
"message": "no access",
|
|
"sqlstate": None,
|
|
})
|
|
assert exc.sqlstate == "KAMC1"
|
|
|
|
|
|
def test_rest_exception_pickle_with_sqlstate():
|
|
e1 = RestException({"error_code": "PERMISSION_DENIED", "message": "no access"})
|
|
e2 = pickle.loads(pickle.dumps(e1))
|
|
assert e1.error_code == e2.error_code
|
|
assert e1.message == e2.message
|
|
assert e1.sqlstate == e2.sqlstate
|
|
assert e1.error_class == e2.error_class
|
|
|
|
|
|
def test_rest_exception_unrecognized_error_code():
|
|
exc = RestException({"error_code": "weird error", "messages": "something"})
|
|
# Unrecognized error codes fall back to INTERNAL_ERROR, which maps to XXMC0
|
|
assert exc.sqlstate == "XXMC0"
|
|
assert exc.error_class == "CP_INTERNAL_ERROR"
|