GOV-667: Fix CI and minor bugs

Author: ankitpatnaik-atlan

pyatlan/client/aio/client.py

@@ -168,7 +168,9 @@ def __init__(self, **kwargs):
 
         # Create async session with custom transport that supports retry and proxy
         self._async_session = httpx.AsyncClient(
-            transport=PyatlanAsyncTransport(retry=self.retry, client=self, **transport_kwargs),
+            transport=PyatlanAsyncTransport(
+                retry=self.retry, client=self, **transport_kwargs
+            ),
             headers={
                 "x-atlan-agent": "sdk",
                 "x-atlan-agent-id": "python",
@@ -980,7 +982,9 @@ async def max_retries(  # type: ignore[override,misc]
         if self.verify is not None:
             transport_kwargs["verify"] = self.verify
 
-        new_transport = PyatlanAsyncTransport(retry=max_retries, client=self, **transport_kwargs)
+        new_transport = PyatlanAsyncTransport(
+            retry=max_retries, client=self, **transport_kwargs
+        )
         session._transport = new_transport
 
         LOGGER.debug(

pyatlan/client/common/transport.py

@@ -6,6 +6,7 @@
 Provides duplicate AuthPolicy detection logic used by both
 PyatlanSyncTransport and PyatlanAsyncTransport.
 """
+
 from __future__ import annotations
 
 import json
@@ -21,7 +22,9 @@
 logger = logging.getLogger(__name__)
 
 
-def build_policy_search_request(policy_name: str, persona_guid: str) -> IndexSearchRequest:
+def build_policy_search_request(
+    policy_name: str, persona_guid: str
+) -> IndexSearchRequest:
     """Build an IndexSearchRequest to find an existing AuthPolicy by name and persona."""
     query = Bool(
         filter=[
@@ -36,7 +39,9 @@ def build_policy_search_request(policy_name: str, persona_guid: str) -> IndexSea
     )
 
 
-def create_mock_response(existing_policy: dict, temp_guid: str = "-1") -> httpx.Response:
+def create_mock_response(
+    existing_policy: dict, temp_guid: str = "-1"
+) -> httpx.Response:
     """Build a mock bulk-entity response containing an already-created policy."""
     response_body = {
         "mutatedEntities": {"CREATE": [existing_policy]},
@@ -59,7 +64,14 @@ def parse_auth_policy_entity(request: httpx.Request) -> Optional[tuple[str, str,
     if not request.content:
         return None
 
-    body = json.loads(request.content.decode("utf-8"))
+    try:
+        body = json.loads(request.content.decode("utf-8"))
+    except (json.JSONDecodeError, UnicodeDecodeError):
+        logger.debug(
+            "parse_auth_policy_entity: failed to decode request body, skipping duplicate check"
+        )
+        return None
+
     for entity in body.get("entities", []):
         if entity.get("typeName") != "AuthPolicy":
             continue
@@ -73,7 +85,9 @@ def parse_auth_policy_entity(request: httpx.Request) -> Optional[tuple[str, str,
     return None
 
 
-def find_existing_policy(client: Any, policy_name: str, persona_guid: str) -> Optional[dict]:
+def find_existing_policy(
+    client: Any, policy_name: str, persona_guid: str
+) -> Optional[dict]:
     """
     Search for an existing AuthPolicy by name and persona GUID (synchronous).
 
@@ -89,7 +103,7 @@ def find_existing_policy(client: Any, policy_name: str, persona_guid: str) -> Op
     except Exception as e:
         raise ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY.exception_with_parameters(
             policy_name, persona_guid, str(e)
-        )
+        ) from e
 
 
 async def find_existing_policy_async(
@@ -110,7 +124,7 @@ async def find_existing_policy_async(
     except Exception as e:
         raise ErrorCode.UNABLE_TO_SEARCH_EXISTING_POLICY.exception_with_parameters(
             policy_name, persona_guid, str(e)
-        )
+        ) from e
 
 
 def check_for_duplicate_policy(
@@ -159,7 +173,9 @@ async def check_for_duplicate_policy_async(
         return None
 
     policy_name, persona_guid, temp_guid = parsed
-    existing_policy = await find_existing_policy_async(client, policy_name, persona_guid)
+    existing_policy = await find_existing_policy_async(
+        client, policy_name, persona_guid
+    )
     if existing_policy:
         logger.info(
             f"Found existing policy '{policy_name}' with guid "

pyatlan/client/transport.py

@@ -1,4 +1,3 @@
-# type: ignore
 """
 Custom HTTP transport with retry support for Atlan Python SDK.
 
@@ -131,6 +130,7 @@ def _retry_operation(
                 response = e
                 continue
 
+            assert isinstance(response, httpx.Response)
             if retry.is_exhausted() or not retry.is_retryable_status_code(
                 response.status_code
             ):
@@ -250,6 +250,7 @@ async def _retry_operation_async(
                 response = e
                 continue
 
+            assert isinstance(response, httpx.Response)
             if retry.is_exhausted() or not retry.is_retryable_status_code(
                 response.status_code
             ):

tests/integration/transport_test.py

@@ -9,6 +9,7 @@
 IndexSearch duplicate-check runs against the real server, validating
 the full transport logic.
 """
+
 from unittest.mock import patch
 
 import httpx
@@ -102,7 +103,7 @@ def intercepting_handle(request: httpx.Request) -> httpx.Response:
             return _build_fake_bulk_response(policy_name, persona.guid)
         return original_inner_handle(request)
 
-    transport._transport.handle_request = intercepting_handle
+    transport._transport.handle_request = intercepting_handle  # type: ignore[method-assign]
 
     try:
         policy = Persona.create_metadata_policy(
@@ -123,7 +124,7 @@ def intercepting_handle(request: httpx.Request) -> httpx.Response:
         )
     finally:
         client._session._transport = original_transport
-        transport._transport.handle_request = original_inner_handle
+        transport._transport.handle_request = original_inner_handle  # type: ignore[method-assign]
 
 
 def test_duplicate_prevention_short_circuits_when_policy_exists(
@@ -173,7 +174,7 @@ def intercepting_handle(request: httpx.Request) -> httpx.Response:
             return _build_fake_bulk_response(policy_name, persona.guid)
         return original_inner_handle(request)
 
-    transport._transport.handle_request = intercepting_handle
+    transport._transport.handle_request = intercepting_handle  # type: ignore[method-assign]
 
     try:
         # Patch find_existing_policy in the common transport module so the
@@ -204,4 +205,4 @@ def intercepting_handle(request: httpx.Request) -> httpx.Response:
         )
     finally:
         client._session._transport = original_transport
-        transport._transport.handle_request = original_inner_handle
+        transport._transport.handle_request = original_inner_handle  # type: ignore[method-assign]

tests/unit/test_transport.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2025 Atlan Pte. Ltd.
 """Unit tests for pyatlan.client.transport and pyatlan.client.common.transport."""
-from unittest.mock import AsyncMock, MagicMock, patch
+
+from unittest.mock import AsyncMock, MagicMock
 
 import httpx
 import pytest
@@ -17,7 +18,6 @@
 )
 from pyatlan.client.constants import BULK_UPDATE
 from pyatlan.client.transport import PyatlanAsyncTransport, PyatlanSyncTransport
-from pyatlan.errors import ErrorCode
 
 
 # ---------------------------------------------------------------------------
@@ -111,6 +111,14 @@ def test_uses_custom_temp_guid(self):
         policy_name, persona_guid, temp_guid = parse_auth_policy_entity(req)
         assert temp_guid == "-3"
 
+    def test_returns_none_for_invalid_json(self):
+        req = httpx.Request("POST", BULK_URL, content=b"{not valid json")
+        assert parse_auth_policy_entity(req) is None
+
+    def test_returns_none_for_invalid_utf8(self):
+        req = httpx.Request("POST", BULK_URL, content=b"\xff\xfe")
+        assert parse_auth_policy_entity(req) is None
+
 
 # ---------------------------------------------------------------------------
 # create_mock_response
@@ -400,9 +408,7 @@ def _make_transport(self, client=None):
     @pytest.mark.asyncio
     async def test_duplicate_prevention_short_circuits_retry(self):
         mock_client = MagicMock()
-        mock_client._call_api = AsyncMock(
-            return_value={"entities": [EXISTING_POLICY]}
-        )
+        mock_client._call_api = AsyncMock(return_value={"entities": [EXISTING_POLICY]})
         transport = self._make_transport(client=mock_client)
 
         call_count = 0