Skip to content

Commit f983456

Browse files
yuokadayuokada
committed
Rewrite test codes to replace httpretty with responses
1 parent 296c7fd commit f983456

File tree

2 files changed

+73
-72
lines changed

2 files changed

+73
-72
lines changed

‎tests/unit/oauth_test_utils.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import uuid
1515
from collections import namedtuple
1616

17-
import httpretty
17+
import responses
1818

1919
from trino import constants
2020

@@ -91,14 +91,14 @@ def __call__(self, request, uri, response_headers):
9191

9292
def _get_token_requests(challenge_id):
9393
return list(filter(
94-
lambda r: r.method == "GET" and r.path == f"/{TOKEN_PATH}/{challenge_id}",
95-
httpretty.latest_requests()))
94+
lambda r: r.method == "GET" and r.url == f"{TOKEN_RESOURCE}/{challenge_id}",
95+
responses.calls))
9696

9797

9898
def _post_statement_requests():
9999
return list(filter(
100-
lambda r: r.method == "POST" and r.path == constants.URL_STATEMENT_PATH,
101-
httpretty.latest_requests()))
100+
lambda r: r.method == "POST" and r.url == f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
101+
responses.calls))
102102

103103

104104
class MultithreadedTokenServer:
@@ -111,15 +111,15 @@ def __init__(self, sample_post_response_data, attempts=1):
111111
self.attempts = attempts
112112

113113
# bind post statement
114-
httpretty.register_uri(
115-
method=httpretty.POST,
116-
uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
114+
responses.add(
115+
method=responses.POST,
116+
url=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
117117
body=self.post_statement_callback)
118118

119119
# bind get token
120-
httpretty.register_uri(
121-
method=httpretty.GET,
122-
uri=re.compile(rf"{TOKEN_RESOURCE}/.*"),
120+
responses.add(
121+
method=responses.GET,
122+
url=re.compile(rf"{TOKEN_RESOURCE}/.*"),
123123
body=self.get_token_callback)
124124

125125
# noinspection PyUnusedLocal

‎tests/unit/test_client.py‎

Lines changed: 62 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -23,15 +23,14 @@
2323
from zoneinfo import ZoneInfoNotFoundError
2424

2525
import gssapi
26-
import httpretty
2726
import keyring
2827
try:
2928
import orjson as json
3029
except ImportError:
3130
import json
3231
import pytest
3332
import requests
34-
from httpretty import httprettified
33+
import responses
3534
from requests_gssapi.exceptions import SPNEGOExchangeError
3635
from requests_kerberos.exceptions import KerberosExchangeError
3736
from tzlocal import get_localzone_name # type: ignore
@@ -352,9 +351,9 @@ def long_call(request, uri, headers):
352351
time.sleep(timeout * 2)
353352
return (200, headers, "delayed success")
354353

355-
httpretty.enable()
356-
for method in [httpretty.POST, httpretty.GET]:
357-
httpretty.register_uri(method, url, body=long_call)
354+
responses.start()
355+
for method in [responses.POST, responses.GET]:
356+
responses.add_callback(method, url, callback=long_call)
358357

359358
# timeout without retry
360359
for request_timeout in [timeout, (timeout, timeout)]:
@@ -375,12 +374,12 @@ def long_call(request, uri, headers):
375374
with pytest.raises(requests.exceptions.Timeout):
376375
req.post("select 1")
377376

378-
httpretty.disable()
379-
httpretty.reset()
377+
responses.stop()
378+
responses.reset()
380379

381380

382381
@pytest.mark.parametrize("attempts", [1, 3, 5])
383-
@httprettified
382+
@responses.activate
384383
def test_oauth2_authentication_flow(attempts, sample_post_response_data):
385384
token = str(uuid.uuid4())
386385
challenge_id = str(uuid.uuid4())
@@ -391,17 +390,17 @@ def test_oauth2_authentication_flow(attempts, sample_post_response_data):
391390
post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
392391

393392
# bind post statement
394-
httpretty.register_uri(
395-
method=httpretty.POST,
396-
uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
397-
body=post_statement_callback)
393+
responses.add_callback(
394+
method=responses.POST,
395+
url=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
396+
callback=post_statement_callback)
398397

399398
# bind get token
400399
get_token_callback = GetTokenCallback(token_server, token, attempts)
401-
httpretty.register_uri(
402-
method=httpretty.GET,
403-
uri=token_server,
404-
body=get_token_callback)
400+
responses.add_callback(
401+
method=responses.GET,
402+
url=token_server,
403+
callback=get_token_callback)
405404

406405
redirect_handler = RedirectHandler()
407406

@@ -422,7 +421,7 @@ def test_oauth2_authentication_flow(attempts, sample_post_response_data):
422421
assert len(_get_token_requests(challenge_id)) == attempts
423422

424423

425-
@httprettified
424+
@responses.activate
426425
def test_oauth2_refresh_token_flow(sample_post_response_data):
427426
token = str(uuid.uuid4())
428427
challenge_id = str(uuid.uuid4())
@@ -432,17 +431,17 @@ def test_oauth2_refresh_token_flow(sample_post_response_data):
432431
post_statement_callback = PostStatementCallback(None, token_server, [token], sample_post_response_data)
433432

434433
# bind post statement
435-
httpretty.register_uri(
436-
method=httpretty.POST,
437-
uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
438-
body=post_statement_callback)
434+
responses.add_callback(
435+
method=responses.POST,
436+
url=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
437+
callback=post_statement_callback)
439438

440439
# bind get token
441440
get_token_callback = GetTokenCallback(token_server, token)
442-
httpretty.register_uri(
443-
method=httpretty.GET,
444-
uri=token_server,
445-
body=get_token_callback)
441+
responses.add_callback(
442+
method=responses.GET,
443+
url=token_server,
444+
callback=get_token_callback)
446445

447446
redirect_handler = RedirectHandlerWithException(
448447
trino.exceptions.TrinoAuthError(
@@ -465,7 +464,7 @@ def test_oauth2_refresh_token_flow(sample_post_response_data):
465464

466465

467466
@pytest.mark.parametrize("attempts", [6, 10])
468-
@httprettified
467+
@responses.activate
469468
def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data):
470469
token = str(uuid.uuid4())
471470
challenge_id = str(uuid.uuid4())
@@ -476,17 +475,17 @@ def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data):
476475
post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
477476

478477
# bind post statement
479-
httpretty.register_uri(
480-
method=httpretty.POST,
481-
uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
482-
body=post_statement_callback)
478+
responses.add_callback(
479+
method=responses.POST,
480+
url=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
481+
callback=post_statement_callback)
483482

484483
# bind get token
485484
get_token_callback = GetTokenCallback(token_server, token, attempts)
486-
httpretty.register_uri(
487-
method=httpretty.GET,
488-
uri=f"{TOKEN_RESOURCE}/{challenge_id}",
489-
body=get_token_callback)
485+
responses.add_callback(
486+
method=responses.GET,
487+
url=f"{TOKEN_RESOURCE}/{challenge_id}",
488+
callback=get_token_callback)
490489

491490
redirect_handler = RedirectHandler()
492491

@@ -514,13 +513,13 @@ def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data):
514513
('x_redirect_server="redirect_server", x_token_server="token_server"', 'Error: header info didn\'t match x_redirect_server="redirect_server", x_token_server="token_server"'), # noqa: E501
515514
('Bearer x_redirect_server="redirect_server"', 'Error: header info didn\'t have x_token_server'),
516515
])
517-
@httprettified
516+
@responses.activate
518517
def test_oauth2_authentication_missing_headers(header, error):
519518
# bind post statement
520-
httpretty.register_uri(
521-
method=httpretty.POST,
522-
uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
523-
adding_headers={'WWW-Authenticate': header},
519+
responses.add(
520+
method=responses.POST,
521+
url=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
522+
headers={'WWW-Authenticate': header},
524523
status=401)
525524

526525
request = TrinoRequest(
@@ -548,7 +547,7 @@ def test_oauth2_authentication_missing_headers(header, error):
548547
'x_token_server="{token_server}"'
549548
'Bearer x_redirect_server="{redirect_server}",x_token_server="{token_server}",additional_challenge',
550549
])
551-
@httprettified
550+
@responses.activate
552551
def test_oauth2_header_parsing(header, sample_post_response_data):
553552
token = str(uuid.uuid4())
554553
challenge_id = str(uuid.uuid4())
@@ -565,17 +564,17 @@ def post_statement(request, uri, response_headers):
565564
'Basic realm': '"Trino"'}, ""]
566565

567566
# bind post statement
568-
httpretty.register_uri(
569-
method=httpretty.POST,
570-
uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
571-
body=post_statement)
567+
responses.add_callback(
568+
method=responses.POST,
569+
url=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
570+
callback=post_statement)
572571

573572
# bind get token
574573
get_token_callback = GetTokenCallback(token_server, token)
575-
httpretty.register_uri(
576-
method=httpretty.GET,
577-
uri=token_server,
578-
body=get_token_callback)
574+
responses.add_callback(
575+
method=responses.GET,
576+
url=token_server,
577+
callback=get_token_callback)
579578

580579
redirect_handler = RedirectHandler()
581580

@@ -596,8 +595,7 @@ def post_statement(request, uri, response_headers):
596595
assert len(_get_token_requests(challenge_id)) == 1
597596

598597

599-
@pytest.mark.parametrize("http_status", [400, 401, 500])
600-
@httprettified
598+
@responses.activate
601599
def test_oauth2_authentication_fail_token_server(http_status, sample_post_response_data):
602600
token = str(uuid.uuid4())
603601
challenge_id = str(uuid.uuid4())
@@ -608,16 +606,18 @@ def test_oauth2_authentication_fail_token_server(http_status, sample_post_respon
608606
post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
609607

610608
# bind post statement
611-
httpretty.register_uri(
612-
method=httpretty.POST,
613-
uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
614-
body=post_statement_callback)
615-
616-
httpretty.register_uri(
617-
method=httpretty.GET,
618-
uri=f"{TOKEN_RESOURCE}/{challenge_id}",
609+
responses.add(
610+
method=responses.POST,
611+
url=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
612+
json=post_statement_callback,
613+
)
614+
615+
responses.add(
616+
method=responses.GET,
617+
url=f"{TOKEN_RESOURCE}/{challenge_id}",
619618
status=http_status,
620-
body="error")
619+
body="error",
620+
)
621621

622622
redirect_handler = RedirectHandler()
623623

@@ -628,7 +628,8 @@ def test_oauth2_authentication_fail_token_server(http_status, sample_post_respon
628628
user="test",
629629
),
630630
http_scheme=constants.HTTPS,
631-
auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler))
631+
auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler),
632+
)
632633

633634
with pytest.raises(trino.exceptions.TrinoAuthError) as exp:
634635
request.post("select 1")
@@ -639,7 +640,7 @@ def test_oauth2_authentication_fail_token_server(http_status, sample_post_respon
639640
assert len(_get_token_requests(challenge_id)) == 1
640641

641642

642-
@httprettified
643+
@responses.activate
643644
def test_multithreaded_oauth2_authentication_flow(sample_post_response_data):
644645
redirect_handler = RedirectHandler()
645646
auth = trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler)

0 commit comments

Comments
 (0)