def test_call_endpoints_raises_exceptions():
with mock.patch("mlflow.utils.rest_utils.call_endpoint") as mock_call_endpoint:
response_proto = GetRun.Response()
mock_call_endpoint.side_effect = [
RestException({"error_code": ErrorCode.Name(ENDPOINT_NOT_FOUND)}),
RestException({"error_code": ErrorCode.Name(ENDPOINT_NOT_FOUND)}),
]
host_only = MlflowHostCreds("http://my-host")
endpoints = [("/my/endpoint", "POST"), ("/my/endpoint", "GET")]
with pytest.raises(RestException):
call_endpoints(host_only, endpoints, "", response_proto)
mock_call_endpoint.side_effect = [RestException({}), None]
with pytest.raises(RestException):
call_endpoints(host_only, endpoints, "", response_proto)
def _call_endpoint(self, api, json_body, call_all_endpoints=False):
response_proto = api.Response()
if call_all_endpoints:
endpoints = _METHOD_TO_ALL_INFO[api]
return call_endpoints(self.get_host_creds(), endpoints, json_body,
response_proto)
else:
endpoint, method = _METHOD_TO_INFO[api]
return call_endpoint(self.get_host_creds(), endpoint, method,
json_body, response_proto)
def test_call_endpoints():
with mock.patch(
"mlflow.utils.rest_utils.call_endpoint") as mock_call_endpoint:
response_proto = GetRun.Response()
mock_call_endpoint.side_effect = [
RestException({"error_code": ErrorCode.Name(ENDPOINT_NOT_FOUND)}),
None,
]
host_only = MlflowHostCreds("http://my-host")
endpoints = [("/my/endpoint", "POST"), ("/my/endpoint", "GET")]
resp = call_endpoints(host_only, endpoints, "", response_proto)
mock_call_endpoint.assert_has_calls([
mock.call(host_only, endpoint, method, "", response_proto)
for endpoint, method in endpoints
])
assert resp is None