Skip to content

Commit 868d3c5

Browse files
sfc-gh-mkubiksfc-gh-pczajka
sfc-gh-mkubik
authored and
sfc-gh-pczajka
committed
SNOW-1944162 Add tests for programmatic access token (#2183)
1 parent c40d89b commit 868d3c5

File tree

5 files changed

+257
-0
lines changed

5 files changed

+257
-0
lines changed

‎src/snowflake/connector/connection.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1141,6 +1141,8 @@ def __open_connection(self):
11411141
backoff_generator=self._backoff_generator,
11421142
)
11431143
elif self._authenticator == PROGRAMMATIC_ACCESS_TOKEN:
1144+
if not self._token and self._password:
1145+
self._token = self._password
11441146
self.auth_class = AuthByPAT(self._token)
11451147
else:
11461148
# okta URL, e.g., https://<account>.okta.com/

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
{
2+
"mappings": [
3+
{
4+
"scenarioName": "Invalid PAT authentication flow",
5+
"requiredScenarioState": "Started",
6+
"newScenarioState": "Authentication failed",
7+
"request": {
8+
"urlPathPattern": "/session/v1/login-request.*",
9+
"method": "POST",
10+
"bodyPatterns": [
11+
{
12+
"equalToJson" : {
13+
"data": {
14+
"LOGIN_NAME": "testUser",
15+
"AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
16+
"TOKEN": "some PAT"
17+
}
18+
},
19+
"ignoreExtraElements" : true
20+
}
21+
]
22+
},
23+
"response": {
24+
"status": 200,
25+
"jsonBody": {
26+
"data": {
27+
"nextAction": "RETRY_LOGIN",
28+
"authnMethod": "PAT",
29+
"signInOptions": {}
30+
},
31+
"code": "394400",
32+
"message": "Programmatic access token is invalid.",
33+
"success": false,
34+
"headers": null
35+
}
36+
}
37+
}
38+
]
39+
}

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
{
2+
"mappings": [
3+
{
4+
"scenarioName": "Successful PAT authentication flow",
5+
"requiredScenarioState": "Started",
6+
"newScenarioState": "Authenticated",
7+
"request": {
8+
"urlPathPattern": "/session/v1/login-request.*",
9+
"method": "POST",
10+
"bodyPatterns": [
11+
{
12+
"equalToJson" : {
13+
"data": {
14+
"LOGIN_NAME": "testUser",
15+
"AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
16+
"TOKEN": "some PAT"
17+
}
18+
},
19+
"ignoreExtraElements" : true
20+
},
21+
{
22+
"matchesJsonPath": {
23+
"expression": "$.data.PASSWORD",
24+
"absent": "(absent)"
25+
}
26+
}
27+
]
28+
},
29+
"response": {
30+
"status": 200,
31+
"jsonBody": {
32+
"data": {
33+
"masterToken": "master token",
34+
"token": "session token",
35+
"validityInSeconds": 3600,
36+
"masterValidityInSeconds": 14400,
37+
"displayUserName": "OAUTH_TEST_AUTH_CODE",
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/data/wiremock/mappings/generic/snowflake_disconnect_successful.json

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
{
2+
"requiredScenarioState": "Connected",
3+
"newScenarioState": "Disconnected",
4+
"request": {
5+
"urlPathPattern": "/session",
6+
"method": "POST",
7+
"queryParameters": {
8+
"delete": {
9+
"matches": "true"
10+
}
11+
}
12+
},
13+
"response": {
14+
"status": 200,
15+
"jsonBody": {
16+
"code": 200,
17+
"message": "done",
18+
"success": true
19+
}
20+
}
21+
}

‎test/unit/test_programmatic_access_token.py

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

0 commit comments

Comments
 (0)