Skip to content

Commit ce78617

Browse files
sfc-gh-pczajkasfc-gh-pczajka
committed
[Async] #2183 add async version of the test
1 parent 868d3c5 commit ce78617

File tree

7 files changed

+328
-2
lines changed

7 files changed

+328
-2
lines changed

‎src/snowflake/connector/aio/_connection.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,6 +303,8 @@ async def __open_connection(self):
303303
backoff_generator=self._backoff_generator,
304304
)
305305
elif self._authenticator == PROGRAMMATIC_ACCESS_TOKEN:
306+
if not self._token and self._password:
307+
self._token = self._password
306308
self.auth_class = AuthByPAT(self._token)
307309
elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
308310
self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (

‎test/data/wiremock/mappings/auth/pat/invalid_token.json

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,14 @@
99
"method": "POST",
1010
"bodyPatterns": [
1111
{
12-
"equalToJson" : {
12+
"equalToJson": {
1313
"data": {
1414
"LOGIN_NAME": "testUser",
1515
"AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
1616
"TOKEN": "some PAT"
1717
}
1818
},
19-
"ignoreExtraElements" : true
19+
"ignoreExtraElements": true
2020
}
2121
]
2222
},

‎test/data/wiremock/mappings/auth/pat/invalid_token_async.json

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
{
2+
"mappings": [
3+
{
4+
"request": {
5+
"urlPathPattern": "/session/v1/login-request.*",
6+
"method": "POST",
7+
"bodyPatterns": [
8+
{
9+
"matchesJsonPath": {
10+
"expression": "$.data.CLIENT_ENVIRONMENT.APPLICATION",
11+
"equalTo": "AsyncioPythonConnector"
12+
}
13+
},
14+
{
15+
"equalToJson": {
16+
"data": {
17+
"LOGIN_NAME": "testUser",
18+
"AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
19+
"TOKEN": "some PAT"
20+
}
21+
},
22+
"ignoreExtraElements": true
23+
}
24+
]
25+
},
26+
"response": {
27+
"status": 200,
28+
"headers": {
29+
"Content-Type": "application/json"
30+
},
31+
"jsonBody": {
32+
"data": {
33+
"nextAction": "RETRY_LOGIN",
34+
"authnMethod": "PAT",
35+
"signInOptions": {}
36+
},
37+
"code": "394400",
38+
"message": "Programmatic access token is invalid.",
39+
"success": false,
40+
"headers": null
41+
}
42+
}
43+
}
44+
]
45+
}

‎test/data/wiremock/mappings/auth/pat/successful_flow_async.json

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
{
2+
"mappings": [
3+
{
4+
"request": {
5+
"urlPathPattern": "/session/v1/login-request.*",
6+
"method": "POST",
7+
"bodyPatterns": [
8+
{
9+
"matchesJsonPath": {
10+
"expression": "$.data.CLIENT_ENVIRONMENT.APPLICATION",
11+
"equalTo": "AsyncioPythonConnector"
12+
}
13+
},
14+
{
15+
"equalToJson": {
16+
"data": {
17+
"LOGIN_NAME": "testUser",
18+
"AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
19+
"TOKEN": "some PAT"
20+
}
21+
},
22+
"ignoreExtraElements": true
23+
},
24+
{
25+
"matchesJsonPath": {
26+
"expression": "$.data.PASSWORD",
27+
"absent": "(absent)"
28+
}
29+
}
30+
]
31+
},
32+
"response": {
33+
"status": 200,
34+
"headers": {
35+
"Content-Type": "application/json"
36+
},
37+
"jsonBody": {
38+
"data": {
39+
"masterToken": "master token",
40+
"token": "session token",
41+
"validityInSeconds": 3600,
42+
"masterValidityInSeconds": 14400,
43+
"displayUserName": "testUser",
44+
"serverVersion": "8.48.0 b2024121104444034239f05",
45+
"firstLogin": false,
46+
"remMeToken": null,
47+
"remMeValidityInSeconds": 0,
48+
"healthCheckInterval": 45,
49+
"newClientForUpgrade": "3.12.3",
50+
"sessionId": 1172562260498,
51+
"parameters": [
52+
{
53+
"name": "CLIENT_PREFETCH_THREADS",
54+
"value": 4
55+
}
56+
],
57+
"sessionInfo": {
58+
"databaseName": "TEST_DB",
59+
"schemaName": "TEST_JDBC",
60+
"warehouseName": "TEST_XSMALL",
61+
"roleName": "ANALYST"
62+
},
63+
"idToken": null,
64+
"idTokenValidityInSeconds": 0,
65+
"responseData": null,
66+
"mfaToken": null,
67+
"mfaTokenValidityInSeconds": 0
68+
},
69+
"code": null,
70+
"message": null,
71+
"success": true
72+
}
73+
}
74+
}
75+
]
76+
}

‎test/data/wiremock/mappings/auth/pat/successful_flow_password_async.json

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
{
2+
"mappings": [
3+
{
4+
"request": {
5+
"urlPathPattern": "/session/v1/login-request.*",
6+
"method": "POST",
7+
"bodyPatterns": [
8+
{
9+
"matchesJsonPath": {
10+
"expression": "$.data.CLIENT_ENVIRONMENT.APPLICATION",
11+
"equalTo": "AsyncioPythonConnector"
12+
}
13+
},
14+
{
15+
"equalToJson": {
16+
"data": {
17+
"LOGIN_NAME": "testUser",
18+
"AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
19+
"TOKEN": "some PAT"
20+
}
21+
},
22+
"ignoreExtraElements": true
23+
}
24+
]
25+
},
26+
"response": {
27+
"status": 200,
28+
"headers": {
29+
"Content-Type": "application/json"
30+
},
31+
"jsonBody": {
32+
"data": {
33+
"masterToken": "master token",
34+
"token": "session token",
35+
"validityInSeconds": 3600,
36+
"masterValidityInSeconds": 14400,
37+
"displayUserName": "testUser",
38+
"serverVersion": "8.48.0 b2024121104444034239f05",
39+
"firstLogin": false,
40+
"remMeToken": null,
41+
"remMeValidityInSeconds": 0,
42+
"healthCheckInterval": 45,
43+
"newClientForUpgrade": "3.12.3",
44+
"sessionId": 1172562260498,
45+
"parameters": [
46+
{
47+
"name": "CLIENT_PREFETCH_THREADS",
48+
"value": 4
49+
}
50+
],
51+
"sessionInfo": {
52+
"databaseName": "TEST_DB",
53+
"schemaName": "TEST_JDBC",
54+
"warehouseName": "TEST_XSMALL",
55+
"roleName": "ANALYST"
56+
},
57+
"idToken": null,
58+
"idTokenValidityInSeconds": 0,
59+
"responseData": null,
60+
"mfaToken": null,
61+
"mfaTokenValidityInSeconds": 0
62+
},
63+
"code": null,
64+
"message": null,
65+
"success": true
66+
}
67+
}
68+
}
69+
]
70+
}

‎test/unit/aio/__init__.py

Whitespace-only changes.

‎test/unit/aio/test_programmatic_access_token_async.py

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
#
2+
# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
3+
#
4+
5+
from __future__ import annotations
6+
7+
import pathlib
8+
from typing import Any, Generator
9+
10+
import pytest
11+
12+
try:
13+
from snowflake.connector.aio import SnowflakeConnection
14+
from snowflake.connector.network import PROGRAMMATIC_ACCESS_TOKEN
15+
except ImportError:
16+
pass
17+
18+
import snowflake.connector.errors
19+
20+
from ...wiremock.wiremock_utils import WiremockClient
21+
22+
23+
@pytest.fixture(scope="session")
24+
def wiremock_client() -> Generator[WiremockClient | Any, Any, None]:
25+
with WiremockClient() as client:
26+
yield client
27+
28+
29+
@pytest.mark.skipolddriver
30+
@pytest.mark.asyncio
31+
async def test_valid_pat_async(wiremock_client: WiremockClient) -> None:
32+
wiremock_data_dir = (
33+
pathlib.Path(__file__).parent.parent.parent
34+
/ "data"
35+
/ "wiremock"
36+
/ "mappings"
37+
/ "auth"
38+
/ "pat"
39+
)
40+
41+
wiremock_generic_data_dir = (
42+
pathlib.Path(__file__).parent.parent.parent
43+
/ "data"
44+
/ "wiremock"
45+
/ "mappings"
46+
/ "generic"
47+
)
48+
49+
wiremock_client.import_mapping(wiremock_data_dir / "successful_flow_async.json")
50+
wiremock_client.add_mapping(
51+
wiremock_generic_data_dir / "snowflake_disconnect_successful.json"
52+
)
53+
54+
connection = SnowflakeConnection(
55+
user="testUser",
56+
authenticator=PROGRAMMATIC_ACCESS_TOKEN,
57+
token="some PAT",
58+
account="testAccount",
59+
protocol="http",
60+
host=wiremock_client.wiremock_host,
61+
port=wiremock_client.wiremock_http_port,
62+
)
63+
await connection.connect()
64+
await connection.close()
65+
66+
67+
@pytest.mark.skipolddriver
68+
@pytest.mark.asyncio
69+
async def test_invalid_pat_async(wiremock_client: WiremockClient) -> None:
70+
wiremock_data_dir = (
71+
pathlib.Path(__file__).parent.parent.parent
72+
/ "data"
73+
/ "wiremock"
74+
/ "mappings"
75+
/ "auth"
76+
/ "pat"
77+
)
78+
wiremock_client.import_mapping(wiremock_data_dir / "invalid_token_async.json")
79+
80+
with pytest.raises(snowflake.connector.errors.DatabaseError) as execinfo:
81+
connection = SnowflakeConnection(
82+
user="testUser",
83+
authenticator=PROGRAMMATIC_ACCESS_TOKEN,
84+
token="some PAT",
85+
account="testAccount",
86+
protocol="http",
87+
host=wiremock_client.wiremock_host,
88+
port=wiremock_client.wiremock_http_port,
89+
)
90+
await connection.connect()
91+
92+
assert str(execinfo.value).endswith("Programmatic access token is invalid.")
93+
94+
95+
@pytest.mark.skipolddriver
96+
@pytest.mark.asyncio
97+
async def test_pat_as_password_async(wiremock_client: WiremockClient) -> None:
98+
wiremock_data_dir = (
99+
pathlib.Path(__file__).parent.parent.parent
100+
/ "data"
101+
/ "wiremock"
102+
/ "mappings"
103+
/ "auth"
104+
/ "pat"
105+
)
106+
107+
wiremock_generic_data_dir = (
108+
pathlib.Path(__file__).parent.parent.parent
109+
/ "data"
110+
/ "wiremock"
111+
/ "mappings"
112+
/ "generic"
113+
)
114+
115+
wiremock_client.import_mapping(
116+
wiremock_data_dir / "successful_flow_password_async.json"
117+
)
118+
wiremock_client.add_mapping(
119+
wiremock_generic_data_dir / "snowflake_disconnect_successful.json"
120+
)
121+
122+
connection = SnowflakeConnection(
123+
user="testUser",
124+
authenticator=PROGRAMMATIC_ACCESS_TOKEN,
125+
token=None,
126+
password="some PAT",
127+
account="testAccount",
128+
protocol="http",
129+
host=wiremock_client.wiremock_host,
130+
port=wiremock_client.wiremock_http_port,
131+
)
132+
await connection.connect()
133+
await connection.close()

0 commit comments

Comments
 (0)