fix to pass test

Author: roger-zhangg

aws_lambda_powertools/tracing/provider/xray_tracer.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from contextlib import contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from numbers import Number
 from typing import Generator, Sequence
 
@@ -49,7 +49,7 @@ def trace(self, name: str) -> Generator[XraySpan, None, None]:
         with self.in_subsegment(name=name) as sub_segment:
             yield XraySpan(subsegment=sub_segment)
 
-    @contextmanager
+    @asynccontextmanager
     async def trace_async(self, name: str) -> Generator[XraySpan, None, None]:
         async with self.in_subsegment_async(name=name) as sub_segment:
             yield XraySpan(subsegment=sub_segment)
@@ -61,6 +61,7 @@ def set_attribute(self, key: str, value: str | Number | bool, **kwargs) -> None:
             self.put_annotation(key=key, value=value)
 
     def patch(self, modules: Sequence[str]) -> None:
+        # defined in init
         pass
 
     def patch_all(self) -> None:

aws_lambda_powertools/tracing/tracer.py

@@ -12,6 +12,7 @@
 from ..shared.lazy_import import LazyLoader
 from ..shared.types import AnyCallableT
 from .base import BaseProvider, BaseSegment
+from .provider import xray_tracer
 
 is_cold_start = True
 logger = logging.getLogger(__name__)
@@ -808,11 +809,8 @@ def _patch_xray_provider(self):
         # Due to Lazy Import, we need to activate `core` attrib via import
         # we also need to include `patch`, `patch_all` methods
         # to ensure patch calls are done via the provider
-        from aws_xray_sdk.core import xray_recorder  # type: ignore
 
-        provider = xray_recorder
-        provider.patch = aws_xray_sdk.core.patch
-        provider.patch_all = aws_xray_sdk.core.patch_all
+        provider = xray_tracer.XrayProvider()
 
         return provider
 
@@ -827,7 +825,7 @@ def _disable_xray_trace_batching(self):
         aws_xray_sdk.core.xray_recorder.configure(streaming_threshold=0)
 
     def _is_xray_provider(self):
-        return "aws_xray_sdk" in self.provider.__module__
+        return isinstance(self.provider, xray_tracer.XrayProvider)
 
     def ignore_endpoint(self, hostname: Optional[str] = None, urls: Optional[List[str]] = None):
         """If you want to ignore certain httplib requests you can do so based on the hostname or URL that is being

tests/unit/test_tracing.py

@@ -33,9 +33,11 @@ def __init__(
             self.put_metadata_mock = put_metadata_mock or mocker.MagicMock()
             self.put_annotation_mock = put_annotation_mock or mocker.MagicMock()
             self.in_subsegment = in_subsegment or mocker.MagicMock()
+            self.trace = self.in_subsegment
             self.patch_mock = patch_mock or mocker.MagicMock()
             self.disable_tracing_provider_mock = disable_tracing_provider_mock or mocker.MagicMock()
             self.in_subsegment_async = in_subsegment_async or mocker.MagicMock(spec=True)
+            self.trace_async = self.in_subsegment_async
 
         def put_metadata(self, *args, **kwargs):
             return self.put_metadata_mock(*args, **kwargs)