rebrandly-otel 0.1.18__tar.gz → 0.1.20__tar.gz

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rebrandly_otel
-Version: 0.1.18
+Version: 0.1.20
 Summary: Python OTEL wrapper by Rebrandly
 Home-page: https://github.com/rebrandly/rebrandly-otel-python
 Author: Antonio Romano
@@ -360,6 +360,93 @@ if __name__ == '__main__':
     app.run(debug=True)
 ```
 
+###
+FastAPI
+
+```python
+
+# main_fastapi.py
+from fastapi import FastAPI, HTTPException, Depends
+from contextlib import asynccontextmanager
+from src.rebrandly_otel import otel, logger, force_flush
+from src.fastapi_support import setup_fastapi, get_current_span
+from datetime import datetime
+from typing import Optional
+import uvicorn
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    logger.info("FastAPI application starting up")
+    yield
+    # Shutdown
+    logger.info("FastAPI application shutting down")
+    force_flush()
+
+app = FastAPI(title="FastAPI OTEL Example", lifespan=lifespan)
+
+# Setup FastAPI with OTEL
+setup_fastapi(otel, app)
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    logger.info("Health check requested")
+    return {"status": "healthy"}
+
+@app.post("/process")
+@app.get("/process")
+async def process(span = Depends(get_current_span)):
+    """Process endpoint with custom span."""
+    with otel.span("process_request"):
+        logger.info("Processing request")
+
+        # You can also use the injected span directly
+        if span:
+            span.add_event("custom_processing_event", {
+                "timestamp": datetime.now().isoformat()
+            })
+
+        # Simulate some processing
+        result = {
+            "processed": True,
+            "timestamp": datetime.now().isoformat()
+        }
+
+        logger.info(f"Returning result: {result}")
+        return result
+
+@app.get("/error")
+async def error():
+    """Endpoint that raises an error."""
+    logger.error("Error endpoint called")
+    raise HTTPException(status_code=400, detail="Simulated error")
+
+@app.get("/exception")
+async def exception():
+    """Endpoint that raises an unhandled exception."""
+    logger.error("Exception endpoint called")
+    raise ValueError("Simulated unhandled exception")
+
+@app.get("/items/{item_id}")
+async def get_item(item_id: int, q: Optional[str] = None):
+    """Example endpoint with path and query parameters."""
+    with otel.span("fetch_item", attributes={"item_id": item_id, "query": q}):
+        logger.info(f"Fetching item {item_id} with query: {q}")
+
+        if item_id == 999:
+            raise HTTPException(status_code=404, detail="Item not found")
+
+        return {
+            "item_id": item_id,
+            "name": f"Item {item_id}",
+            "query": q
+        }
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+```
+
 ### More examples
 You can find More examples [here](examples)
 

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/README.md

@@ -339,6 +339,93 @@ if __name__ == '__main__':
     app.run(debug=True)
 ```
 
+###
+FastAPI
+
+```python
+
+# main_fastapi.py
+from fastapi import FastAPI, HTTPException, Depends
+from contextlib import asynccontextmanager
+from src.rebrandly_otel import otel, logger, force_flush
+from src.fastapi_support import setup_fastapi, get_current_span
+from datetime import datetime
+from typing import Optional
+import uvicorn
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    logger.info("FastAPI application starting up")
+    yield
+    # Shutdown
+    logger.info("FastAPI application shutting down")
+    force_flush()
+
+app = FastAPI(title="FastAPI OTEL Example", lifespan=lifespan)
+
+# Setup FastAPI with OTEL
+setup_fastapi(otel, app)
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    logger.info("Health check requested")
+    return {"status": "healthy"}
+
+@app.post("/process")
+@app.get("/process")
+async def process(span = Depends(get_current_span)):
+    """Process endpoint with custom span."""
+    with otel.span("process_request"):
+        logger.info("Processing request")
+
+        # You can also use the injected span directly
+        if span:
+            span.add_event("custom_processing_event", {
+                "timestamp": datetime.now().isoformat()
+            })
+
+        # Simulate some processing
+        result = {
+            "processed": True,
+            "timestamp": datetime.now().isoformat()
+        }
+
+        logger.info(f"Returning result: {result}")
+        return result
+
+@app.get("/error")
+async def error():
+    """Endpoint that raises an error."""
+    logger.error("Error endpoint called")
+    raise HTTPException(status_code=400, detail="Simulated error")
+
+@app.get("/exception")
+async def exception():
+    """Endpoint that raises an unhandled exception."""
+    logger.error("Exception endpoint called")
+    raise ValueError("Simulated unhandled exception")
+
+@app.get("/items/{item_id}")
+async def get_item(item_id: int, q: Optional[str] = None):
+    """Example endpoint with path and query parameters."""
+    with otel.span("fetch_item", attributes={"item_id": item_id, "query": q}):
+        logger.info(f"Fetching item {item_id} with query: {q}")
+
+        if item_id == 999:
+            raise HTTPException(status_code=404, detail="Item not found")
+
+        return {
+            "item_id": item_id,
+            "name": f"Item {item_id}",
+            "query": q
+        }
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+```
+
 ### More examples
 You can find More examples [here](examples)
 

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/rebrandly_otel.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rebrandly_otel
-Version: 0.1.18
+Version: 0.1.20
 Summary: Python OTEL wrapper by Rebrandly
 Home-page: https://github.com/rebrandly/rebrandly-otel-python
 Author: Antonio Romano
@@ -360,6 +360,93 @@ if __name__ == '__main__':
     app.run(debug=True)
 ```
 
+###
+FastAPI
+
+```python
+
+# main_fastapi.py
+from fastapi import FastAPI, HTTPException, Depends
+from contextlib import asynccontextmanager
+from src.rebrandly_otel import otel, logger, force_flush
+from src.fastapi_support import setup_fastapi, get_current_span
+from datetime import datetime
+from typing import Optional
+import uvicorn
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    logger.info("FastAPI application starting up")
+    yield
+    # Shutdown
+    logger.info("FastAPI application shutting down")
+    force_flush()
+
+app = FastAPI(title="FastAPI OTEL Example", lifespan=lifespan)
+
+# Setup FastAPI with OTEL
+setup_fastapi(otel, app)
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    logger.info("Health check requested")
+    return {"status": "healthy"}
+
+@app.post("/process")
+@app.get("/process")
+async def process(span = Depends(get_current_span)):
+    """Process endpoint with custom span."""
+    with otel.span("process_request"):
+        logger.info("Processing request")
+
+        # You can also use the injected span directly
+        if span:
+            span.add_event("custom_processing_event", {
+                "timestamp": datetime.now().isoformat()
+            })
+
+        # Simulate some processing
+        result = {
+            "processed": True,
+            "timestamp": datetime.now().isoformat()
+        }
+
+        logger.info(f"Returning result: {result}")
+        return result
+
+@app.get("/error")
+async def error():
+    """Endpoint that raises an error."""
+    logger.error("Error endpoint called")
+    raise HTTPException(status_code=400, detail="Simulated error")
+
+@app.get("/exception")
+async def exception():
+    """Endpoint that raises an unhandled exception."""
+    logger.error("Exception endpoint called")
+    raise ValueError("Simulated unhandled exception")
+
+@app.get("/items/{item_id}")
+async def get_item(item_id: int, q: Optional[str] = None):
+    """Example endpoint with path and query parameters."""
+    with otel.span("fetch_item", attributes={"item_id": item_id, "query": q}):
+        logger.info(f"Fetching item {item_id} with query: {q}")
+
+        if item_id == 999:
+            raise HTTPException(status_code=404, detail="Item not found")
+
+        return {
+            "item_id": item_id,
+            "name": f"Item {item_id}",
+            "query": q
+        }
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+```
+
 ### More examples
 You can find More examples [here](examples)
 

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/rebrandly_otel.egg-info/SOURCES.txt

@@ -12,4 +12,5 @@ src/logs.py
 src/metrics.py
 src/otel_utils.py
 src/rebrandly_otel.py
-src/traces.py
+src/traces.py
+tests/test_usage.py

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="rebrandly_otel",
-    version="0.1.18",
+    version="0.1.20",
     author="Antonio Romano",
     author_email="antonio@rebrandly.com",
     description="Python OTEL wrapper by Rebrandly",

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/src/logs.py

@@ -1,6 +1,7 @@
 # logs.py
 """Logging implementation for Rebrandly OTEL SDK."""
 import logging
+import sys
 from typing import Optional
 from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
 from opentelemetry.sdk._logs.export import (
@@ -34,12 +35,15 @@ class RebrandlyLogger:
             self._provider.add_log_record_processor(SimpleLogRecordProcessor(console_exporter))
 
         # Add OTLP exporter if configured
-        if get_otlp_endpoint():
-            otlp_exporter = OTLPLogExporter(endpoint=get_otlp_endpoint())
+        otel_endpoint = get_otlp_endpoint()
+        if otel_endpoint:
+            otlp_exporter = OTLPLogExporter(
+                timeout=5,
+                endpoint=otel_endpoint
+            )
             batch_processor = BatchLogRecordProcessor(otlp_exporter, export_timeout_millis=get_millis_batch_time())
             self._provider.add_log_record_processor(batch_processor)
 
-        # Set as global provider
         set_logger_provider(self._provider)
 
         # Configure standard logging

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/src/metrics.py

@@ -78,8 +78,12 @@ class RebrandlyMeter:
             readers.append(console_reader)
 
         # Add OTLP exporter if configured
-        if get_otlp_endpoint() is not None:
-            otlp_exporter = OTLPMetricExporter(endpoint=get_otlp_endpoint())
+        otel_endpoint = get_otlp_endpoint()
+        if otel_endpoint is not None:
+            otlp_exporter = OTLPMetricExporter(
+                endpoint=otel_endpoint,
+                timeout=5
+            )
             otlp_reader = PeriodicExportingMetricReader(otlp_exporter, export_interval_millis=get_millis_batch_time())
             readers.append(otlp_reader)
 

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/src/otel_utils.py

@@ -3,6 +3,8 @@
 
 import os
 import sys
+import grpc
+import json
 
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.semconv.attributes import service_attributes
@@ -33,8 +35,9 @@ def get_package_version():
         try:
             from importlib_metadata import version, PackageNotFoundError
             return version('rebrandly_otel')
-        except:
-            return '?.0.1'
+        except Exception as e:
+            print(f"[OTEL Utils] Warning: Could not get package version: {e}")
+            return '0.1.0'
 
 
 def get_service_name(service_name: str = None) -> str:
@@ -49,10 +52,31 @@ def get_service_version(service_version: str = None) -> str:
     return service_version
 
 
-def get_otlp_endpoint(otlp_endpoint: str = None) -> str:
-    if otlp_endpoint is None:
-        return os.environ.get('OTEL_EXPORTER_OTLP_ENDPOINT', None)
-    return otlp_endpoint
+def get_otlp_endpoint(otlp_endpoint: str = None) -> str | None:
+    endpoint = otlp_endpoint or os.environ.get('OTEL_EXPORTER_OTLP_ENDPOINT', None)
+
+    if endpoint is not None:
+        try:
+            from urllib.parse import urlparse
+
+            # Parse the endpoint
+            parsed = urlparse(endpoint if '://' in endpoint else f'http://{endpoint}')
+            host = parsed.hostname
+            port = parsed.port
+
+            # Test gRPC connection
+            channel = grpc.insecure_channel(f'{host}:{port}')
+            try:
+                # Wait for the channel to be ready
+                grpc.channel_ready_future(channel).result(timeout=3)
+                return endpoint
+            finally:
+                channel.close()
+
+        except Exception as e:
+            print(f"[OTEL] Failed to connect to OTLP endpoint {endpoint}: {e}")
+            return None
+    return endpoint
 
 def is_otel_debug() -> bool:
     return os.environ.get('OTEL_DEBUG', 'false').lower() == 'true'
@@ -61,5 +85,25 @@ def is_otel_debug() -> bool:
 def get_millis_batch_time():
     try:
         return int(os.environ.get('BATCH_EXPORT_TIME_MILLIS', 100))
-    except:
-        return 5000
+    except Exception as e:
+        print(f"[OTEL Utils] Warning: Invalid BATCH_EXPORT_TIME_MILLIS value, using default 5000ms: {e}")
+        return 5000
+
+def extract_event_from(message) -> str | None:
+    body = None
+    if 'body' in message:
+        body = message['body']
+    if 'Body' in message:
+        body = message['Body']
+    if 'Message' in message:
+        body = message['Message']
+    if 'Sns' in message and 'Message' in message['Sns']:
+        body = message['Sns']['Message']
+    if body is not None:
+        try:
+            jbody = json.loads(body)
+            if 'event' in jbody:
+                return jbody['event']
+        except:
+            pass
+    return None

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/src/rebrandly_otel.py

@@ -12,6 +12,7 @@ from opentelemetry import baggage, propagate, context
 from .traces import RebrandlyTracer
 from .metrics import RebrandlyMeter
 from .logs import RebrandlyLogger
+from .otel_utils import extract_event_from
 
 
 T = TypeVar('T')
@@ -260,8 +261,11 @@ class RebrandlyOTEL:
                 try:
                     # Create span and execute function
                     span_function = self.span
-                    if record is not None and ('MessageAttributes' in record or 'messageAttributes' in record):
+                    if record is not None and (('MessageAttributes' in record or 'messageAttributes' in record) or ('Sns' in record and 'MessageAttributes' in record['Sns'])):
                         span_function = self.aws_message_span
+                        evt = extract_event_from(record)
+                        if evt:
+                            span_attributes['event.type'] = evt
 
                     with span_function(span_name, message=record, attributes=span_attributes, kind=kind) as span_context:
                         # Add processing start event with standardized name
@@ -370,6 +374,7 @@ class RebrandlyOTEL:
             time.sleep(0.1)
 
         except Exception as e:
+            print(f"[Rebrandly OTEL] Error during force flush: {e}")
             success = False
 
         return success
@@ -486,6 +491,10 @@ class RebrandlyOTEL:
             if 'awsRegion' in message:
                 combined_attributes['cloud.region'] = message['awsRegion']
 
+            evt = extract_event_from(message)
+            if evt:
+                combined_attributes['event.type'] = evt
+
 
         # Use the tracer's start_span method directly to ensure it works
         # This creates a child span of whatever is currently active

{rebrandly_otel-0.1.18 → rebrandly_otel-0.1.20}/src/traces.py

@@ -32,8 +32,12 @@ class RebrandlyTracer:
             self._provider.add_span_processor(SimpleSpanProcessor(console_exporter))
 
         # Add OTLP exporter if configured
-        if get_otlp_endpoint() is not None:
-            otlp_exporter = OTLPSpanExporter(endpoint=get_otlp_endpoint())
+        otel_endpoint = get_otlp_endpoint()
+        if otel_endpoint is not None:
+            otlp_exporter = OTLPSpanExporter(
+                endpoint=otel_endpoint,
+                timeout=5
+            )
 
             # Use batch processor for production
             batch_processor = BatchSpanProcessor(otlp_exporter, export_timeout_millis=get_millis_batch_time())
@@ -139,14 +143,14 @@ class RebrandlyTracer:
         if target_span and hasattr(target_span, 'record_exception'):
             if exception is not None:
                 target_span.record_exception(exception)
-            target_span.set_status(trace.Status(trace.StatusCode.ERROR, str(exception)), description=msg)
+                target_span.set_status(trace.Status(trace.StatusCode.ERROR, str(exception)))
 
 
     def record_span_success(self, span: Optional[Span] = None, msg: Optional[str] = None):
-        """Record an exception on a span."""
+        """Record success on a span."""
         target_span = span or self.get_current_span()
-        if target_span and hasattr(target_span, 'record_exception'):
-            target_span.set_status(trace.Status(trace.StatusCode.OK), description=msg)
+        if target_span and hasattr(target_span, 'set_status'):
+            target_span.set_status(trace.Status(trace.StatusCode.OK))
 
 
     def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None, span: Optional[Span] = None):

rebrandly_otel-0.1.20/tests/test_usage.py

@@ -0,0 +1,336 @@
+import unittest
+from unittest.mock import patch, MagicMock, ANY, call
+import importlib
+import os
+
+# Import the module under test
+import src.rebrandly_otel
+from src.rebrandly_otel import RebrandlyOTEL, otel, lambda_handler, aws_message_handler, span, force_flush, shutdown
+from opentelemetry.trace import SpanKind, Status, StatusCode
+
+
+class TestRebrandlyOTELUsage(unittest.TestCase):
+
+    def setUp(self):
+        # Reset the singleton and reload the module before each test
+        RebrandlyOTEL._instance = None
+        RebrandlyOTEL._initialized = False
+        importlib.reload(src.rebrandly_otel)
+
+        # Patch external dependencies at the module level
+        self.mock_propagate = patch('opentelemetry.propagate').start()
+        self.mock_context = patch('opentelemetry.context').start()
+        self.mock_psutil = patch('src.rebrandly_otel.psutil').start()
+
+        # Patch the attributes of the otel instance directly
+        self.mock_tracer = MagicMock()
+        self.mock_meter = MagicMock()
+        self.mock_logger = MagicMock()
+
+        otel._tracer = self.mock_tracer
+        otel._meter = self.mock_meter
+        otel._logger = self.mock_logger
+
+        # Configure mock behaviors for tracer
+        mock_span_instance = MagicMock()
+        mock_span_instance.set_status = MagicMock()
+        mock_span_instance.record_exception = MagicMock()
+        mock_span_instance.set_attribute = MagicMock()
+
+        mock_start_span_context_manager = MagicMock()
+        mock_start_span_context_manager.__enter__.return_value = mock_span_instance
+        mock_start_span_context_manager.__exit__.return_value = None
+
+        self.mock_tracer.start_span.return_value = mock_start_span_context_manager
+
+        # Configure mock behaviors for meter
+        self.mock_meter.GlobalMetrics.memory_usage_bytes.set = MagicMock()
+        self.mock_meter.GlobalMetrics.cpu_usage_percentage.set = MagicMock()
+
+        # Configure mock behaviors for psutil
+        self.mock_psutil.cpu_percent.return_value = 10.0
+        mock_virtual_memory = MagicMock()
+        mock_virtual_memory.used = 1000000
+        mock_virtual_memory.percent = 50.0
+        self.mock_psutil.virtual_memory.return_value = mock_virtual_memory
+
+        self.addCleanup(patch.stopall)
+
+    def test_initialization(self):
+        # otel instance is already created and initialized in setUp
+        self.assertIsInstance(otel, RebrandlyOTEL)
+        # The tracer, meter, logger properties are accessed during otel initialization
+        # which happens when the module is reloaded in setUp.
+        # So, we check that the mocks were assigned.
+        self.assertIs(otel.tracer, self.mock_tracer)
+        self.assertIs(otel.meter, self.mock_meter)
+        self.assertIs(otel.logger, self.mock_logger)
+
+    def test_span_context_manager(self):
+        with span("test_span", attributes={"key": "value"}) as s:
+            self.assertIsNotNone(s)
+            s.set_attribute("another_key", "another_value")
+
+        self.mock_tracer.start_span.assert_called_once_with(
+            "test_span", attributes={"key": "value"}, kind=SpanKind.INTERNAL
+        )
+        span_mock = self.mock_tracer.start_span.return_value.__enter__.return_value
+        self.assertEqual(span_mock.set_status.call_args[0][0].status_code, StatusCode.OK)
+
+    def test_span_context_manager_with_exception(self):
+        with self.assertRaises(ValueError):
+            with span("error_span"):
+                raise ValueError("Something went wrong")
+
+        self.mock_tracer.start_span.assert_called_once_with(
+            "error_span", attributes=None, kind=SpanKind.INTERNAL
+        )
+        span_mock = self.mock_tracer.start_span.return_value.__enter__.return_value
+        span_mock.record_exception.assert_called_once()
+        self.assertEqual(span_mock.set_status.call_args[0][0].status_code, StatusCode.ERROR)
+        self.assertEqual(span_mock.set_status.call_args[0][0].description, "Something went wrong")
+
+    def test_lambda_handler_basic(self):
+        @lambda_handler("my_lambda")
+        def handler(event, context):
+            return {"statusCode": 200, "body": "OK"}
+
+        mock_context = MagicMock()
+        mock_context.aws_request_id = "req123"
+        mock_context.function_arn = "arn:aws:lambda:us-east-1:123456789012:function:my_lambda"
+        mock_context.function_name = "my_lambda"
+        mock_context.function_version = "$LATEST"
+
+        result = handler({}, mock_context)
+
+        self.assertEqual(result, {"statusCode": 200, "body": "OK"})
+        self.mock_tracer.start_span.assert_called_once_with(
+            name="my_lambda",
+            attributes={
+                "faas.trigger": "direct",
+                "faas.execution": "req123",
+                "faas.id": "arn:aws:lambda:us-east-1:123456789012:function:my_lambda",
+                "faas.name": "my_lambda",
+                "faas.version": "$LATEST"
+            },
+            kind=SpanKind.SERVER
+        )
+        span_mock = self.mock_tracer.start_span.return_value.__enter__.return_value
+        self.assertEqual(span_mock.set_status.call_args[0][0].status_code, StatusCode.OK)
+        self.mock_logger.logger.info.assert_called_with(ANY)
+        self.mock_psutil.cpu_percent.assert_called_once()
+        self.mock_psutil.virtual_memory.assert_called_once()
+        self.mock_meter.GlobalMetrics.memory_usage_bytes.set.assert_called_once()
+        self.mock_meter.GlobalMetrics.cpu_usage_percentage.set.assert_called_once()
+        self.mock_tracer.force_flush.assert_called_once()
+        self.mock_meter.force_flush.assert_called_once()
+        self.mock_logger.force_flush.assert_called_once()
+
+    def test_lambda_handler_with_sqs_event(self):
+        @lambda_handler("sqs_processor")
+        def handler(event, context):
+            return {"statusCode": 200}
+
+        sqs_event = {
+            "Records": [
+                {
+                    "messageId": "msg1",
+                    "eventSource": "aws:sqs",
+                    "messageAttributes": {
+                        "traceparent": {"stringValue": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", "dataType": "String"}
+                    }
+                }
+            ]
+        }
+        mock_context = MagicMock()
+        mock_context.aws_request_id = "req456"
+
+        handler(sqs_event, mock_context)
+
+        self.mock_propagate.extract.assert_called_once_with({'traceparent': '00-0123456789abcdef0123456789abcdef-0123456789abcdef-01'})
+        self.mock_context.attach.assert_called_once()
+        self.mock_context.detach.assert_called_once()
+        self.mock_tracer.start_span.assert_called_once_with(
+            name="sqs_processor",
+            attributes=ANY,
+            kind=SpanKind.SERVER
+        )
+        span_mock = self.mock_tracer.start_span.return_value.__enter__.return_value
+        self.assertEqual(span_mock.set_status.call_args[0][0].status_code, StatusCode.OK)
+        self.mock_logger.logger.info.assert_called_with(ANY)
+        self.mock_psutil.cpu_percent.assert_called_once()
+        self.mock_psutil.virtual_memory.assert_called_once()
+        self.mock_meter.GlobalMetrics.memory_usage_bytes.set.assert_called_once()
+        self.mock_meter.GlobalMetrics.cpu_usage_percentage.set.assert_called_once()
+        self.mock_tracer.force_flush.assert_called_once()
+        self.mock_meter.force_flush.assert_called_once()
+        self.mock_logger.force_flush.assert_called_once()
+
+    def test_lambda_handler_with_sns_event(self):
+        @lambda_handler("sns_processor")
+        def handler(event, context):
+            return {"statusCode": 200}
+
+        sns_event = {
+            "Records": [
+                {
+                    "EventSource": "aws:sns",
+                    "Sns": {
+                        "MessageId": "sns_msg1",
+                        "TopicArn": "arn:aws:sns:us-east-1:123456789012:my-topic",
+                        "Subject": "test_subject",
+                        "MessageAttributes": {
+                            "traceparent": {"Type": "String", "Value": "00-fedcba9876543210fedcba9876543210-fedcba9876543210-01"}
+                        }
+                    }
+                }
+            ]
+        }
+        mock_context = MagicMock()
+        mock_context.aws_request_id = "req789"
+
+        handler(sns_event, mock_context)
+
+        self.mock_propagate.extract.assert_called_once_with({'traceparent': '00-fedcba9876543210fedcba9876543210-fedcba9876543210-01'})
+        self.mock_context.attach.assert_called_once()
+        self.mock_context.detach.assert_called_once()
+        self.mock_tracer.start_span.assert_called_once_with(
+            name="sns_processor",
+            attributes=ANY,
+            kind=SpanKind.SERVER
+        )
+        span_mock = self.mock_tracer.start_span.return_value.__enter__.return_value
+        self.assertEqual(span_mock.set_status.call_args[0][0].status_code, StatusCode.OK)
+        self.mock_logger.logger.info.assert_called_with(ANY)
+        self.mock_psutil.cpu_percent.assert_called_once()
+        self.mock_psutil.virtual_memory.assert_called_once()
+        self.mock_meter.GlobalMetrics.memory_usage_bytes.set.assert_called_once()
+        self.mock_meter.GlobalMetrics.cpu_usage_percentage.set.assert_called_once()
+        self.mock_tracer.force_flush.assert_called_once()
+        self.mock_meter.force_flush.assert_called_once()
+        self.mock_logger.force_flush.assert_called_once()
+
+    def test_aws_message_handler_basic(self):
+        @aws_message_handler("msg_processor")
+        def handler(record):
+            return {"processed": True}
+
+        record = {"messageId": "rec123"}
+        result = handler(record)
+
+        self.assertEqual(result, {"processed": True})
+        self.mock_tracer.start_span.assert_called_once_with(
+            name="msg_processor",
+            attributes={'messaging.operation': 'process'},
+            kind=SpanKind.CONSUMER
+        )
+        span_mock = self.mock_tracer.start_span.return_value.__enter__.return_value
+        self.assertEqual(span_mock.set_status.call_args[0][0].status_code, StatusCode.OK)
+        self.mock_psutil.cpu_percent.assert_called_once()
+        self.mock_psutil.virtual_memory.assert_called_once()
+        self.mock_meter.GlobalMetrics.memory_usage_bytes.set.assert_called_once()
+        self.mock_meter.GlobalMetrics.cpu_usage_percentage.set.assert_called_once()
+        self.mock_tracer.force_flush.assert_called_once()
+        self.mock_meter.force_flush.assert_called_once()
+        self.mock_logger.force_flush.assert_called_once()
+
+    def test_aws_message_span_context_manager(self):
+        sqs_message = {
+            "messageId": "sqs_msg_id",
+            "eventSource": "aws:sqs",
+            "awsRegion": "us-east-1",
+            "messageAttributes": {
+                "traceparent": {"stringValue": "00-11111111111111111111111111111111-2222222222222222-01", "dataType": "String"}
+            }
+        }
+        with otel.aws_message_span("process_sqs", message=sqs_message) as s:
+            self.assertIsNotNone(s)
+            s.set_attribute("custom_attr", "custom_value")
+
+        self.mock_tracer.start_span.assert_called_once_with(
+            name="process_sqs",
+            attributes={
+                'messaging.operation': 'process',
+                'messaging.message_id': 'sqs_msg_id',
+                'messaging.system': 'aws:sqs',
+                'cloud.region': 'us-east-1'
+            },
+            kind=SpanKind.CONSUMER
+        )
+        span_mock = self.mock_tracer.start_span.return_value.__enter__.return_value
+        span_mock.set_attribute.assert_called_once_with("custom_attr", "custom_value")
+        self.assertEqual(span_mock.set_status.call_args[0][0].status_code, StatusCode.OK)
+
+    def test_force_flush(self):
+        otel.force_flush()
+        self.mock_tracer.force_flush.assert_called_once()
+        self.mock_meter.force_flush.assert_called_once()
+        self.mock_logger.force_flush.assert_called_once()
+
+    def test_shutdown(self):
+        otel.shutdown()
+        self.mock_tracer.shutdown.assert_called_once()
+        self.mock_meter.shutdown.assert_called_once()
+        self.mock_logger.shutdown.assert_called_once()
+
+    def test_detect_lambda_trigger(self):
+        # Test direct invocation
+        self.assertEqual(otel._detect_lambda_trigger(None), 'direct')
+        self.assertEqual(otel._detect_lambda_trigger({}), 'direct')
+
+        # Test SQS
+        sqs_event = {'Records': [{'eventSource': 'aws:sqs'}]}
+        self.assertEqual(otel._detect_lambda_trigger(sqs_event), 'sqs')
+
+        # Test SNS
+        sns_event = {'Records': [{'eventSource': 'aws:sns'}]}
+        self.assertEqual(otel._detect_lambda_trigger(sns_event), 'sns')
+
+        # Test API Gateway
+        api_gw_event = {'httpMethod': 'GET'}
+        self.assertEqual(otel._detect_lambda_trigger(api_gw_event), 'api_gateway')
+
+        api_gw_v2_event = {'requestContext': {'http': {}}}
+        self.assertEqual(otel._detect_lambda_trigger(api_gw_v2_event), 'api_gateway_v2')
+
+        # Test EventBridge
+        eventbridge_event = {'source': 'aws.events'}
+        self.assertEqual(otel._detect_lambda_trigger(eventbridge_event), 'eventbridge')
+
+        # Test unknown
+        unknown_event = {'some_key': 'some_value'}
+        self.assertEqual(otel._detect_lambda_trigger(unknown_event), 'unknown')
+
+    def test_context_propagation_methods(self):
+        # Mock the global propagate and context objects
+        mock_propagate = MagicMock()
+        mock_context = MagicMock()
+
+        with patch('src.rebrandly_otel.propagate', mock_propagate), patch('src.rebrandly_otel.context', mock_context):
+
+            # Test inject_context
+            carrier = {}
+            otel.inject_context(carrier)
+            mock_propagate.inject.assert_called_once_with(carrier)
+
+            # Test extract_context
+            extracted_ctx = otel.extract_context(carrier)
+            mock_propagate.extract.assert_called_once_with(carrier)
+            self.assertEqual(extracted_ctx, mock_propagate.extract.return_value)
+
+            # Test attach_context
+            mock_propagate.extract.reset_mock()
+            mock_context.attach.reset_mock()
+            token = otel.attach_context(carrier)
+            mock_propagate.extract.assert_called_once_with(carrier)
+            mock_context.attach.assert_called_once_with(mock_propagate.extract.return_value)
+            self.assertEqual(token, mock_context.attach.return_value)
+
+            # Test detach_context
+            mock_context.detach.reset_mock()
+            otel.detach_context(token)
+            mock_context.detach.assert_called_once_with(token)
+
+
+if __name__ == '__main__':
+    unittest.main()