Skip to content

Commit caf84b2

Browse files
npeshkovsfc-gh-turbaszek
npeshkov
authored and
sfc-gh-turbaszek
committed
add unit tests
1 parent b34c4be commit caf84b2

File tree

3 files changed

+189
-3
lines changed

3 files changed

+189
-3
lines changed

‎test/unit/test_auth_callback_server.py‎

Lines changed: 101 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,10 @@ def test_auth_callback_success(monkeypatch, dontwait, timeout, reuse_port) -> No
2424
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
2525
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
2626
test_response: requests.Response | None = None
27-
with AuthHttpServer("http://127.0.0.1/test_request") as callback_server:
27+
with AuthHttpServer(
28+
uri="http://127.0.0.1/test_request",
29+
redirect_uri="http://127.0.0.1/test_request",
30+
) as callback_server:
2831

2932
def request_callback():
3033
nonlocal test_response
@@ -57,7 +60,103 @@ def request_callback():
5760
def test_auth_callback_timeout(monkeypatch, dontwait, timeout, reuse_port) -> None:
5861
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
5962
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
60-
with AuthHttpServer("http://127.0.0.1/test_request") as callback_server:
63+
with AuthHttpServer(
64+
uri="http://127.0.0.1/test_request",
65+
redirect_uri="http://127.0.0.1/test_request",
66+
) as callback_server:
6167
block, client_socket = callback_server.receive_block(timeout=timeout)
6268
assert block is None
6369
assert client_socket is None
70+
71+
72+
@pytest.mark.parametrize(
73+
"socket_host",
74+
[
75+
"127.0.0.1",
76+
"localhost",
77+
],
78+
)
79+
@pytest.mark.parametrize(
80+
"socket_port",
81+
[
82+
"",
83+
":0",
84+
":12345",
85+
],
86+
)
87+
@pytest.mark.parametrize(
88+
"redirect_host",
89+
[
90+
"127.0.0.1",
91+
"localhost",
92+
],
93+
)
94+
@pytest.mark.parametrize(
95+
"redirect_port",
96+
[
97+
"",
98+
":0",
99+
":12345",
100+
],
101+
)
102+
@pytest.mark.parametrize(
103+
"dontwait",
104+
["false", "true"],
105+
)
106+
@pytest.mark.parametrize("reuse_port", ["true", "false"])
107+
def test_auth_callback_server_updates_localhost_redirect_uri_port_to_match_socket_port(
108+
monkeypatch,
109+
socket_host,
110+
socket_port,
111+
redirect_host,
112+
redirect_port,
113+
dontwait,
114+
reuse_port,
115+
) -> None:
116+
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
117+
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
118+
with AuthHttpServer(
119+
uri=f"http://{socket_host}{socket_port}/test_request",
120+
redirect_uri=f"http://{redirect_host}{redirect_port}/test_request",
121+
) as callback_server:
122+
assert callback_server._redirect_uri.port == callback_server.port
123+
124+
125+
@pytest.mark.parametrize(
126+
"socket_host",
127+
[
128+
"127.0.0.1",
129+
"localhost",
130+
],
131+
)
132+
@pytest.mark.parametrize(
133+
"socket_port",
134+
[
135+
"",
136+
":0",
137+
":12345",
138+
],
139+
)
140+
@pytest.mark.parametrize(
141+
"redirect_port",
142+
[
143+
"",
144+
":0",
145+
":12345",
146+
],
147+
)
148+
@pytest.mark.parametrize(
149+
"dontwait",
150+
["false", "true"],
151+
)
152+
@pytest.mark.parametrize("reuse_port", ["true", "false"])
153+
def test_auth_callback_server_does_not_updates_nonlocalhost_redirect_uri_port_to_match_socket_port(
154+
monkeypatch, socket_host, socket_port, redirect_port, dontwait, reuse_port
155+
) -> None:
156+
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
157+
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
158+
redirect_uri = f"http://not_localhost{redirect_port}/test_request"
159+
with AuthHttpServer(
160+
uri=f"http://{socket_host}{socket_port}/test_request", redirect_uri=redirect_uri
161+
) as callback_server:
162+
assert callback_server.redirect_uri == redirect_uri

‎test/unit/test_auth_oauth_auth_code.py‎

Lines changed: 67 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
import unittest.mock as mock
77
from test.helpers import apply_auth_class_update_body, create_mock_auth_body
8-
from unittest.mock import patch
8+
from unittest.mock import PropertyMock, patch
99

1010
import pytest
1111

@@ -285,3 +285,69 @@ def mock_request_tokens(self, **kwargs):
285285
assert isinstance(conn.auth_class, AuthByOauthCode)
286286

287287
conn.close()
288+
289+
290+
@pytest.mark.parametrize("redirect_uri", ["https://redirect/uri"])
291+
@pytest.mark.parametrize("rtr_enabled", [True, False])
292+
def test_auth_oauth_auth_code_uses_redirect_uri(
293+
redirect_uri, rtr_enabled: bool, omit_oauth_urls_check
294+
):
295+
"""Test that the redirect URI is used correctly in the OAuth authorization code flow."""
296+
auth = AuthByOauthCode(
297+
"app",
298+
"clientId",
299+
"clientSecret",
300+
"auth_url",
301+
"tokenRequestUrl",
302+
redirect_uri,
303+
"scope",
304+
"host",
305+
pkce_enabled=False,
306+
enable_single_use_refresh_tokens=rtr_enabled,
307+
)
308+
309+
def fake_get_request_token_response(_, fields: dict[str, str]):
310+
if rtr_enabled:
311+
assert fields.get("enable_single_use_refresh_tokens") == "true"
312+
else:
313+
assert "enable_single_use_refresh_tokens" not in fields
314+
return ("access_token", "refresh_token")
315+
316+
with patch(
317+
"snowflake.connector.auth.AuthByOauthCode._construct_authorization_request",
318+
return_value="authorization_request",
319+
) as mock_construct_authorization_request:
320+
with patch(
321+
"snowflake.connector.auth.AuthByOauthCode._receive_authorization_callback",
322+
return_value=("code", auth._state),
323+
):
324+
with patch(
325+
"snowflake.connector.auth.AuthByOauthCode._ask_authorization_callback_from_user",
326+
return_value=("code", auth._state),
327+
):
328+
with patch(
329+
"snowflake.connector.auth.AuthByOauthCode._get_request_token_response",
330+
side_effect=fake_get_request_token_response,
331+
) as mock_get_request_token_response:
332+
with patch(
333+
"snowflake.connector.auth._http_server.AuthHttpServer.redirect_uri",
334+
return_value=redirect_uri,
335+
new_callable=PropertyMock,
336+
):
337+
auth.prepare(
338+
conn=None,
339+
authenticator=OAUTH_AUTHORIZATION_CODE,
340+
service_name=None,
341+
account="acc",
342+
user="user",
343+
)
344+
mock_construct_authorization_request.assert_called_once_with(
345+
redirect_uri
346+
)
347+
assert mock_get_request_token_response.call_count == 1
348+
assert (
349+
mock_get_request_token_response.call_args[0][1][
350+
"redirect_uri"
351+
]
352+
== redirect_uri
353+
)

‎test/unit/test_oauth_token.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ def test_oauth_code_successful_flow(
124124
omit_oauth_urls_check,
125125
) -> None:
126126
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
127+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
127128

128129
wiremock_client.import_mapping(
129130
wiremock_oauth_authorization_code_dir / "successful_flow.json"
@@ -166,6 +167,7 @@ def test_oauth_code_invalid_state(
166167
omit_oauth_urls_check,
167168
) -> None:
168169
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
170+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
169171

170172
wiremock_client.import_mapping(
171173
wiremock_oauth_authorization_code_dir / "invalid_state_error.json"
@@ -201,6 +203,7 @@ def test_oauth_code_scope_error(
201203
monkeypatch,
202204
) -> None:
203205
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
206+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
204207

205208
wiremock_client.import_mapping(
206209
wiremock_oauth_authorization_code_dir / "invalid_scope_error.json"
@@ -237,6 +240,7 @@ def test_oauth_code_token_request_error(
237240
omit_oauth_urls_check,
238241
) -> None:
239242
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
243+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
240244

241245
with WiremockClient() as wiremock_client:
242246
wiremock_client.import_mapping(
@@ -275,6 +279,7 @@ def test_oauth_code_browser_timeout(
275279
omit_oauth_urls_check,
276280
) -> None:
277281
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
282+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
278283

279284
wiremock_client.import_mapping(
280285
wiremock_oauth_authorization_code_dir
@@ -316,6 +321,7 @@ def test_oauth_code_custom_urls(
316321
omit_oauth_urls_check,
317322
) -> None:
318323
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
324+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
319325

320326
wiremock_client.import_mapping(
321327
wiremock_oauth_authorization_code_dir / "external_idp_custom_urls.json"
@@ -359,6 +365,7 @@ def test_oauth_code_local_application_custom_urls_successful_flow(
359365
omit_oauth_urls_check,
360366
) -> None:
361367
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
368+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
362369

363370
wiremock_client.import_mapping(
364371
wiremock_oauth_authorization_code_dir
@@ -403,6 +410,7 @@ def test_oauth_code_successful_refresh_token_flow(
403410
omit_oauth_urls_check,
404411
) -> None:
405412
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
413+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
406414

407415
wiremock_client.import_mapping(
408416
wiremock_generic_mappings_dir / "snowflake_login_failed.json"
@@ -463,6 +471,7 @@ def test_oauth_code_expired_refresh_token_flow(
463471
omit_oauth_urls_check,
464472
) -> None:
465473
monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
474+
monkeypatch.setenv("SNOWFLAKE_OAUTH_SOCKET_PORT", "8009")
466475

467476
wiremock_client.import_mapping(
468477
wiremock_generic_mappings_dir / "snowflake_login_failed.json"
@@ -546,6 +555,10 @@ def test_client_creds_successful_flow(
546555
monkeypatch,
547556
temp_cache,
548557
) -> None:
558+
monkeypatch.setenv(
559+
"SNOWFLAKE_OAUTH_SOCKET_PORT", wiremock_client.wiremock_http_port
560+
)
561+
549562
wiremock_client.import_mapping(
550563
wiremock_oauth_client_creds_dir / "successful_flow.json"
551564
)
@@ -596,6 +609,10 @@ def test_client_creds_token_request_error(
596609
wiremock_generic_mappings_dir,
597610
monkeypatch,
598611
) -> None:
612+
monkeypatch.setenv(
613+
"SNOWFLAKE_OAUTH_SOCKET_PORT", wiremock_client.wiremock_http_port
614+
)
615+
599616
wiremock_client.import_mapping(
600617
wiremock_oauth_client_creds_dir / "token_request_error.json"
601618
)
@@ -637,6 +654,10 @@ def test_client_creds_expired_refresh_token_flow(
637654
monkeypatch,
638655
temp_cache,
639656
) -> None:
657+
monkeypatch.setenv(
658+
"SNOWFLAKE_OAUTH_SOCKET_PORT", wiremock_client.wiremock_http_port
659+
)
660+
640661
wiremock_client.import_mapping(
641662
wiremock_generic_mappings_dir / "snowflake_login_failed.json"
642663
)

0 commit comments

Comments
 (0)