def test_authentication_fail_retry(monkeypatch):
post_retry = RetryRecorder(error=KerberosExchangeError())
monkeypatch.setattr(TrinoRequest.http.Session, "post", post_retry)
get_retry = RetryRecorder(error=KerberosExchangeError())
monkeypatch.setattr(TrinoRequest.http.Session, "get", get_retry)
attempts = 3
kerberos_auth = KerberosAuthentication()
req = TrinoRequest(
host="coordinator",
port=8080,
user="******",
http_scheme=constants.HTTPS,
auth=kerberos_auth,
max_attempts=attempts,
)
with pytest.raises(KerberosExchangeError):
req.post("URL")
assert post_retry.retry_count == attempts
with pytest.raises(KerberosExchangeError):
req.get("URL")
assert post_retry.retry_count == attempts
def test_request_timeout():
timeout = 0.1
http_scheme = "http"
host = "coordinator"
port = 8080
url = http_scheme + "://" + host + ":" + str(
port) + constants.URL_STATEMENT_PATH
def long_call(request, uri, headers):
time.sleep(timeout * 2)
return (200, headers, "delayed success")
httpretty.enable()
for method in [httpretty.POST, httpretty.GET]:
httpretty.register_uri(method, url, body=long_call)
# timeout without retry
for request_timeout in [timeout, (timeout, timeout)]:
req = TrinoRequest(
host=host,
port=port,
user="******",
http_scheme=http_scheme,
max_attempts=1,
request_timeout=request_timeout,
)
with pytest.raises(requests.exceptions.Timeout):
req.get(url)
with pytest.raises(requests.exceptions.Timeout):
req.post("select 1")
httpretty.disable()
httpretty.reset()
def test_request_headers(monkeypatch):
post_recorder = ArgumentsRecorder()
monkeypatch.setattr(TrinoRequest.http.Session, "post", post_recorder)
get_recorder = ArgumentsRecorder()
monkeypatch.setattr(TrinoRequest.http.Session, "get", get_recorder)
catalog = "test_catalog"
schema = "test_schema"
user = "******"
source = "test_source"
accept_encoding_header = "accept-encoding"
accept_encoding_value = "identity,deflate,gzip"
client_info_header = constants.HEADER_CLIENT_INFO
client_info_value = "some_client_info"
req = TrinoRequest(
host="coordinator",
port=8080,
user=user,
source=source,
catalog=catalog,
schema=schema,
http_scheme="http",
session_properties={},
http_headers={
accept_encoding_header: accept_encoding_value,
client_info_header: client_info_value,
},
redirect_handler=None,
)
def assert_headers(headers):
assert headers[constants.HEADER_CATALOG] == catalog
assert headers[constants.HEADER_SCHEMA] == schema
assert headers[constants.HEADER_SOURCE] == source
assert headers[constants.HEADER_USER] == user
assert headers[constants.HEADER_SESSION] == ""
assert headers[accept_encoding_header] == accept_encoding_value
assert headers[client_info_header] == client_info_value
assert len(headers.keys()) == 8
req.post("URL")
assert_headers(post_recorder.kwargs["headers"])
req.get("URL")
assert_headers(get_recorder.kwargs["headers"])
def test_request_headers(mock_get_and_post):
get, post = mock_get_and_post
catalog = "test_catalog"
schema = "test_schema"
user = "******"
source = "test_source"
accept_encoding_header = "accept-encoding"
accept_encoding_value = "identity,deflate,gzip"
client_info_header = constants.HEADER_CLIENT_INFO
client_info_value = "some_client_info"
req = TrinoRequest(
host="coordinator",
port=8080,
user=user,
source=source,
catalog=catalog,
schema=schema,
http_scheme="http",
session_properties={},
http_headers={
accept_encoding_header: accept_encoding_value,
client_info_header: client_info_value,
},
redirect_handler=None,
)
def assert_headers(headers):
assert headers[constants.HEADER_CATALOG] == catalog
assert headers[constants.HEADER_SCHEMA] == schema
assert headers[constants.HEADER_SOURCE] == source
assert headers[constants.HEADER_USER] == user
assert headers[constants.HEADER_SESSION] == ""
assert headers[accept_encoding_header] == accept_encoding_value
assert headers[client_info_header] == client_info_value
assert len(headers.keys()) == 8
req.post("URL")
_, post_kwargs = post.call_args
assert_headers(post_kwargs["headers"])
req.get("URL")
_, get_kwargs = get.call_args
assert_headers(get_kwargs["headers"])
def test_503_error_retry(monkeypatch):
http_resp = TrinoRequest.http.Response()
http_resp.status_code = 503
post_retry = RetryRecorder(result=http_resp)
monkeypatch.setattr(TrinoRequest.http.Session, "post", post_retry)
get_retry = RetryRecorder(result=http_resp)
monkeypatch.setattr(TrinoRequest.http.Session, "get", get_retry)
attempts = 3
req = TrinoRequest(host="coordinator",
port=8080,
user="******",
max_attempts=attempts)
req.post("URL")
assert post_retry.retry_count == attempts
req.get("URL")
assert post_retry.retry_count == attempts
def test_request_session_properties_headers(mock_get_and_post):
get, post = mock_get_and_post
req = TrinoRequest(host="coordinator",
port=8080,
user="******",
session_properties={
"a": "1",
"b": "2",
"c": "more=v1,v2"
})
def assert_headers(headers):
assert headers[constants.HEADER_SESSION] == "a=1,b=2,c=more%3Dv1%2Cv2"
req.post("URL")
_, post_kwargs = post.call_args
assert_headers(post_kwargs["headers"])
req.get("URL")
_, get_kwargs = get.call_args
assert_headers(get_kwargs["headers"])
