Add tests for programmatic access token

DESCRIPTION.md

@@ -15,9 +15,10 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Fixed a bug where privatelink OCSP Cache url could not be determined if privatelink account name was specified in uppercase.
   - Added support for iceberg tables to `write_pandas`.
   - Fixed base64 encoded private key tests.
-  - Added Wiremock tests.
+  - Added Wiremock.
   - Fixed a bug where file permission check happened on Windows
   - Added support for File types.
+  - Added Wiremock tests for programmatic access token.
 
 - v3.13.2(January 29, 2025)
   - Changed not to use scoped temporary objects.

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

@@ -0,0 +1,39 @@
+{
+  "mappings": [
+    {
+      "scenarioName": "Invalid PAT authentication flow",
+      "requiredScenarioState": "Started",
+      "newScenarioState": "Authentication failed",
+      "request": {
+        "urlPathPattern": "/session/v1/login-request.*",
+        "method": "POST",
+        "bodyPatterns": [
+          {
+            "equalToJson" : {
+              "data": {
+                "LOGIN_NAME": "testUser",
+                "AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
+                "TOKEN": "some PAT"
+              }
+            },
+            "ignoreExtraElements" : true
+          }
+        ]
+      },
+      "response": {
+        "status": 200,
+        "jsonBody": {
+          "data": {
+            "nextAction": "RETRY_LOGIN",
+            "authnMethod": "PAT",
+            "signInOptions": {}
+          },
+          "code": "394400",
+          "message": "Programmatic access token is invalid.",
+          "success": false,
+          "headers": null
+        }
+      }
+    }
+  ]
+}

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

@@ -0,0 +1,70 @@
+{
+  "mappings": [
+    {
+      "scenarioName": "Successful PAT authentication flow",
+      "requiredScenarioState": "Started",
+      "newScenarioState": "Authenticated",
+      "request": {
+        "urlPathPattern": "/session/v1/login-request.*",
+        "method": "POST",
+        "bodyPatterns": [
+          {
+            "equalToJson" : {
+              "data": {
+                "LOGIN_NAME": "testUser",
+                "AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN",
+                "TOKEN": "some PAT"
+              }
+            },
+            "ignoreExtraElements" : true
+          },
+          {
+            "matchesJsonPath": {
+              "expression": "$.data.PASSWORD",
+              "absent": "(absent)"
+            }
+          }
+        ]
+      },
+      "response": {
+        "status": 200,
+        "jsonBody": {
+          "data": {
+            "masterToken": "master token",
+            "token": "session token",
+            "validityInSeconds": 3600,
+            "masterValidityInSeconds": 14400,
+            "displayUserName": "OAUTH_TEST_AUTH_CODE",
+            "serverVersion": "8.48.0 b2024121104444034239f05",
+            "firstLogin": false,
+            "remMeToken": null,
+            "remMeValidityInSeconds": 0,
+            "healthCheckInterval": 45,
+            "newClientForUpgrade": "3.12.3",
+            "sessionId": 1172562260498,
+            "parameters": [
+              {
+                "name": "CLIENT_PREFETCH_THREADS",
+                "value": 4
+              }
+            ],
+            "sessionInfo": {
+              "databaseName": "TEST_DB",
+              "schemaName": "TEST_JDBC",
+              "warehouseName": "TEST_XSMALL",
+              "roleName": "ANALYST"
+            },
+            "idToken": null,
+            "idTokenValidityInSeconds": 0,
+            "responseData": null,
+            "mfaToken": null,
+            "mfaTokenValidityInSeconds": 0
+          },
+          "code": null,
+          "message": null,
+          "success": true
+        }
+      }
+    }
+  ]
+}

test/data/wiremock/mappings/generic/snowflake_disconnect_successful.json

@@ -0,0 +1,21 @@
+{
+  "requiredScenarioState": "Connected",
+  "newScenarioState": "Disconnected",
+  "request": {
+    "urlPathPattern": "/session",
+    "method": "POST",
+    "queryParameters": {
+      "delete": {
+        "matches": "true"
+      }
+    }
+  },
+  "response": {
+    "status": 200,
+    "jsonBody": {
+      "code": 200,
+      "message": "done",
+      "success": true
+    }
+  }
+}

test/unit/test_programmatic_access_token.py

@@ -0,0 +1,91 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+import pathlib
+from typing import Any, Generator, Union
+
+import pytest
+
+try:
+    import snowflake.connector
+    from src.snowflake.connector.network import PROGRAMMATIC_ACCESS_TOKEN
+    from src.snowflake.connector.test_util import RUNNING_ON_JENKINS
+except ImportError:
+    import os
+
+    RUNNING_ON_JENKINS = os.getenv("JENKINS_HOME") is not None
+
+from ..wiremock.wiremock_utils import WiremockClient
+
+
+@pytest.fixture(scope="session")
+def wiremock_client() -> Generator[Union[WiremockClient, Any], Any, None]:
+    with WiremockClient() as client:
+        yield client
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.skipif(RUNNING_ON_JENKINS, reason="jenkins doesn't support wiremock tests")
+def test_valid_pat(wiremock_client: WiremockClient) -> None:
+    wiremock_data_dir = (
+        pathlib.Path(__file__).parent.parent
+        / "data"
+        / "wiremock"
+        / "mappings"
+        / "auth"
+        / "pat"
+    )
+
+    wiremock_generic_data_dir = (
+        pathlib.Path(__file__).parent.parent
+        / "data"
+        / "wiremock"
+        / "mappings"
+        / "generic"
+    )
+
+    wiremock_client.import_mapping(wiremock_data_dir / "successful_flow.json")
+    wiremock_client.add_mapping(
+        wiremock_generic_data_dir / "snowflake_disconnect_successful.json"
+    )
+
+    cnx = snowflake.connector.connect(
+        user="testUser",
+        authenticator=PROGRAMMATIC_ACCESS_TOKEN,
+        token="some PAT",
+        account="testAccount",
+        protocol="http",
+        host=wiremock_client.wiremock_host,
+        port=wiremock_client.wiremock_http_port,
+    )
+
+    assert cnx, "invalid cnx"
+    cnx.close()
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.skipif(RUNNING_ON_JENKINS, reason="jenkins doesn't support wiremock tests")
+def test_invalid_pat(wiremock_client: WiremockClient) -> None:
+    wiremock_data_dir = (
+        pathlib.Path(__file__).parent.parent
+        / "data"
+        / "wiremock"
+        / "mappings"
+        / "auth"
+        / "pat"
+    )
+    wiremock_client.import_mapping(wiremock_data_dir / "invalid_token.json")
+
+    with pytest.raises(snowflake.connector.errors.DatabaseError) as execinfo:
+        snowflake.connector.connect(
+            user="testUser",
+            authenticator=PROGRAMMATIC_ACCESS_TOKEN,
+            token="some PAT",
+            account="testAccount",
+            protocol="http",
+            host=wiremock_client.wiremock_host,
+            port=wiremock_client.wiremock_http_port,
+        )
+
+    assert str(execinfo.value).endswith("Programmatic access token is invalid.")