Applied minor fixes and updated test cases

Author: webarch-ai

lib/idp_common_pkg/idp_common/monitoring/models.py

@@ -246,11 +246,18 @@ class MonitoringKPIs:
     active_config_count: int = 0
 
     def compute_derived(self) -> None:
-        """Recompute failure_rate and avg_cost_per_document from raw counts."""
+        """Recompute failure_rate and avg_cost_per_document from raw counts.
+
+        avg_cost_per_document is always (re-)assigned inside the
+        ``total_documents > 0`` branch so that a subsequent call with
+        ``total_cost = 0`` correctly resets it to 0.0 rather than leaving
+        a stale non-zero value from a previous computation.
+        """
         if self.total_documents > 0:
             self.failure_rate = self.failed_documents / self.total_documents
-            if self.total_cost > 0:
-                self.avg_cost_per_document = self.total_cost / self.total_documents
+            self.avg_cost_per_document = (
+                self.total_cost / self.total_documents if self.total_cost > 0 else 0.0
+            )
         else:
             self.failure_rate = 0.0
             self.avg_cost_per_document = 0.0

lib/idp_common_pkg/idp_common/monitoring/settings_cache.py

@@ -204,3 +204,22 @@ def get_setting(key: str, default: str = "") -> str:
 def get_cloudwatch_log_groups() -> List[str]:
     """Return the CloudWatch log group list from the shared default cache."""
     return _default_cache.get_cloudwatch_log_groups()
+
+
+def reset_default_cache(
+    ttl_seconds: int = 300,
+    ssm_client: Optional[Any] = None,
+) -> None:
+    """
+    Replace the module-level singleton with a fresh :class:`SettingsCache`.
+
+    **For testing only.**  Call this in test fixtures or ``teardown`` to
+    prevent state leakage between test cases that exercise the module-level
+    helpers (:func:`get_setting`, :func:`get_cloudwatch_log_groups`).
+
+    Args:
+        ttl_seconds: TTL for the new cache (default: 300).
+        ssm_client:  Optional pre-built SSM mock client to inject.
+    """
+    global _default_cache
+    _default_cache = SettingsCache(ttl_seconds=ttl_seconds, ssm_client=ssm_client)

lib/idp_common_pkg/idp_common/monitoring/stack_utils.py

@@ -34,7 +34,7 @@
 logger = logging.getLogger(__name__)
 
 # ---------------------------------------------------------------------------
-# Module-level lazy boto3 client cache (M-1: avoids creating a new client per call)
+# Module-level lazy boto3 client cache
 # ---------------------------------------------------------------------------
 _cf_client: Optional[Any] = None
 
@@ -79,16 +79,25 @@ def extract_stack_name_from_arn(execution_arn: str) -> str:
 
     try:
         # ARN format: arn:aws:states:<region>:<account>:execution:<state-machine-name>:<exec-id>
-        parts = execution_arn.split(":")
-        if len(parts) < 8:
+        # A valid Step Functions execution ARN has at least 8 colon-separated fields.
+        # Use rsplit(maxsplit=2) so that execution IDs containing colons are handled
+        # correctly (the state-machine name is always the second-to-last component).
+        if execution_arn.count(":") < 7:
+            # Too few fields — definitely not a valid execution ARN
+            return ""
+        parts = execution_arn.rsplit(":", maxsplit=2)
+        # Expected: ["arn:aws:states:...:execution", "<state-machine-name>", "<exec-id>"]
+        if len(parts) < 3:
             return ""
 
-        state_machine_name: str = parts[6]
+        state_machine_name: str = parts[-2]
 
-        # Strip known workflow suffixes to recover the stack name
+        # Strip known workflow suffixes to recover the stack name.
+        # Use endswith() so that a suffix appearing in the middle of the name
+        # is not incorrectly stripped.
         for suffix in _STATE_MACHINE_SUFFIXES:
-            if suffix in state_machine_name:
-                return state_machine_name.replace(suffix, "")
+            if state_machine_name.endswith(suffix):
+                return state_machine_name[: -len(suffix)]
 
         # Fallback: remove the last hyphen-delimited word
         # e.g. "MYSTACK-SomeNewSuffix" → "MYSTACK"

lib/idp_common_pkg/idp_common/monitoring/stepfunctions_service.py

@@ -38,17 +38,18 @@
 logger = logging.getLogger(__name__)
 
 # ---------------------------------------------------------------------------
-# Module-level lazy boto3 client cache (M-1: avoids creating a new client per call)
+# Module-level lazy boto3 client cache, keyed by region so that callers passing
+# an explicit region always get the correct client, even if a default-region
+# client was already initialised first.
 # ---------------------------------------------------------------------------
-_sf_client: Optional[Any] = None
+_sf_clients: Dict[Optional[str], Any] = {}
 
 
 def _get_sf_client(region: Optional[str] = None) -> Any:
-    """Return (and lazily create) a module-level Step Functions boto3 client."""
-    global _sf_client
-    if _sf_client is None:
-        _sf_client = boto3.client("stepfunctions", region_name=region)
-    return _sf_client
+    """Return (and lazily create) a per-region Step Functions boto3 client."""
+    if region not in _sf_clients:
+        _sf_clients[region] = boto3.client("stepfunctions", region_name=region)
+    return _sf_clients[region]
 
 
 # ---------------------------------------------------------------------------
@@ -171,7 +172,10 @@ def get_execution_data(
         result["events"] = events
 
     except sf.exceptions.ExecutionDoesNotExist:
+        # Use a distinct status so callers can tell "not found" from a
+        # transient API error (which leaves status as "UNKNOWN").
         logger.warning("Execution not found: %s", execution_arn)
+        result["status"] = "NOT_FOUND"
     except Exception as exc:  # noqa: BLE001
         logger.warning("Failed to get execution data for %s: %s", execution_arn, exc)
 
@@ -214,7 +218,16 @@ def analyze_execution_timeline(
         }``
     """
     exec_data = get_execution_data(execution_arn, region=region)
-    events: List[Dict[str, Any]] = exec_data.get("events", [])[:max_events]
+    all_events: List[Dict[str, Any]] = exec_data.get("events", [])
+    if len(all_events) > max_events:
+        logger.warning(
+            "Execution %s has %d events; truncating to %d for timeline analysis. "
+            "Increase max_events to avoid missing failure details.",
+            execution_arn,
+            len(all_events),
+            max_events,
+        )
+    events = all_events[:max_events]
 
     timeline: Dict[str, Any] = {
         "execution_arn": execution_arn,
@@ -236,38 +249,51 @@ def analyze_execution_timeline(
             state_name = event.get("stateEnteredEventDetails", {}).get("name", "")
             if state_name:
                 state_starts[state_name] = timestamp
+            continue
 
-        elif event_type in (
-            "TaskStateExited",
-            "TaskFailed",
-            "TaskTimedOut",
-            "ExecutionFailed",
-        ):
+        # --- State-completion and failure events that produce a timeline entry ---
+        if event_type == "TaskStateExited":
+            state_name = event.get("stateExitedEventDetails", {}).get("name", "")
+            is_failure = False
+            status = "SUCCEEDED"
+
+        elif event_type in ("TaskFailed", "TaskTimedOut"):
+            # Task-level failures do not carry stateExitedEventDetails — look
+            # back to the most recently entered state to get the name.
+            state_name = list(state_starts.keys())[-1] if state_starts else "Unknown"
+            is_failure = True
+            status = "FAILED"
+
+        elif event_type == "ExecutionFailed":
             state_name = event.get("stateExitedEventDetails", {}).get("name", "")
-            if not state_name and event_type == "ExecutionFailed":
+            if not state_name:
                 # Attribute failure to the last known entered state
                 state_name = (
                     list(state_starts.keys())[-1] if state_starts else "Unknown"
                 )
+            is_failure = True
+            status = "FAILED"
 
-            is_failure = "Failed" in event_type or "TimedOut" in event_type
-            status = "FAILED" if is_failure else "SUCCEEDED"
-            start = state_starts.get(state_name, "")
-            duration_ms = _compute_duration_ms(start, timestamp)
+        else:
+            # Ignore all other event types (ExecutionStarted, LambdaScheduled, etc.)
+            continue
 
-            states.append(
-                {
-                    "name": state_name,
-                    "status": status,
-                    "start_time": start,
-                    "end_time": timestamp,
-                    "duration_ms": round(duration_ms, 1),
-                    "is_failure": is_failure,
-                }
-            )
+        start = state_starts.get(state_name, "")
+        duration_ms = _compute_duration_ms(start, timestamp)
+
+        states.append(
+            {
+                "name": state_name,
+                "status": status,
+                "start_time": start,
+                "end_time": timestamp,
+                "duration_ms": round(duration_ms, 1),
+                "is_failure": is_failure,
+            }
+        )
 
-            if is_failure and not timeline["failed_state"]:
-                timeline["failed_state"] = state_name
+        if is_failure and not timeline["failed_state"]:
+            timeline["failed_state"] = state_name
 
     timeline["states"] = states
     timeline["failure_details"] = extract_failure_details(events)

lib/idp_common_pkg/idp_common/monitoring/xray_service.py

@@ -43,9 +43,10 @@
 _DEFAULT_XRAY_LOOKBACK_HOURS: int = 24
 
 # ---------------------------------------------------------------------------
-# Module-level lazy boto3 client cache (M-1: avoids creating a new client per call)
+# Module-level lazy boto3 client/resource cache
 # ---------------------------------------------------------------------------
 _xray_client: Optional[Any] = None
+_dynamodb_resource: Optional[Any] = None
 
 
 def _get_xray_client() -> Any:
@@ -56,6 +57,21 @@ def _get_xray_client() -> Any:
     return _xray_client
 
 
+def _get_dynamodb_resource() -> Any:
+    """Return (and lazily create) a module-level DynamoDB boto3 resource."""
+    global _dynamodb_resource
+    if _dynamodb_resource is None:
+        _dynamodb_resource = boto3.resource("dynamodb")
+    return _dynamodb_resource
+
+
+# Named constant for the tracking table sort key to avoid a magic string
+# literal scattered through the code.  The tracking table uses a single-item
+# partition design where the document metadata record is stored with a sentinel
+# sort key value.
+_TRACKING_TABLE_SK: str = "none"
+
+
 # ---------------------------------------------------------------------------
 # Public API
 # ---------------------------------------------------------------------------
@@ -258,9 +274,11 @@ def _get_trace_id_from_dynamodb(document_id: str) -> Optional[str]:
         return None
 
     try:
-        dynamodb = boto3.resource("dynamodb")
+        dynamodb = _get_dynamodb_resource()
         table = dynamodb.Table(table_name)  # type: ignore[attr-defined]
-        response = table.get_item(Key={"PK": f"doc#{document_id}", "SK": "none"})
+        response = table.get_item(
+            Key={"PK": f"doc#{document_id}", "SK": _TRACKING_TABLE_SK}
+        )
         if "Item" in response:
             trace_id = response["Item"].get("TraceId")
             if trace_id:
@@ -292,7 +310,7 @@ def _get_trace_id_from_xray_annotations(
         )
         traces = response.get("TraceSummaries", [])
         if traces:
-            # M-3: Sort by ResponseTime descending to get the most recent trace
+            # Sort by ResponseTime descending to get the most recent trace
             # (covers the case where a document has been reprocessed/retried)
             traces_sorted = sorted(
                 traces,

lib/idp_common_pkg/tests/unit/monitoring/test_stack_utils.py

@@ -19,7 +19,7 @@
 
 
 # ---------------------------------------------------------------------------
-# Reset module-level boto3 client cache between tests (M-1 fix)
+# Reset module-level boto3 client cache between tests
 # ---------------------------------------------------------------------------
 @pytest.fixture(autouse=True)
 def reset_cf_client():

lib/idp_common_pkg/tests/unit/monitoring/test_stepfunctions_service.py

@@ -19,14 +19,15 @@
 
 
 # ---------------------------------------------------------------------------
-# Reset module-level boto3 client cache between tests (M-1 fix)
+# Reset module-level boto3 client cache between tests.
+# _sf_clients is a dict keyed by region; clear it between tests.
 # ---------------------------------------------------------------------------
 @pytest.fixture(autouse=True)
 def reset_sf_client():
-    """Reset the module-level _sf_client singleton before every test."""
-    sf_module._sf_client = None
+    """Reset the module-level _sf_clients dict before every test."""
+    sf_module._sf_clients.clear()
     yield
-    sf_module._sf_client = None
+    sf_module._sf_clients.clear()
 
 
 # ---------------------------------------------------------------------------
@@ -242,6 +243,22 @@ def test_handles_missing_execution_gracefully(self):
             "not found"
         )
 
+        with patch("idp_common.monitoring.stepfunctions_service.boto3") as mock_boto3:
+            mock_boto3.client.return_value = sf
+            result = get_execution_data(_EXEC_ARN)
+
+        # ExecutionDoesNotExist → NOT_FOUND (not UNKNOWN, which is for API errors)
+        assert result["status"] == "NOT_FOUND"
+        assert result["events"] == []
+
+    def test_generic_api_error_returns_unknown_status(self):
+        """A transient API error must leave status as UNKNOWN, not NOT_FOUND."""
+        sf = MagicMock()
+        sf.exceptions.ExecutionDoesNotExist = type(
+            "ExecutionDoesNotExist", (Exception,), {}
+        )
+        sf.describe_execution.side_effect = Exception("Service unavailable")
+
         with patch("idp_common.monitoring.stepfunctions_service.boto3") as mock_boto3:
             mock_boto3.client.return_value = sf
             result = get_execution_data(_EXEC_ARN)
@@ -301,6 +318,38 @@ def test_empty_events_returns_empty_states(self):
         assert timeline["states"] == []
         assert timeline["failed_state"] == ""
 
+    def test_task_failed_event_records_state_name_from_prior_entered(self):
+        """
+        TaskFailed does not carry stateExitedEventDetails, so the state
+        name must be inferred from the last TaskStateEntered event.
+        """
+        events = _make_events(include_failure=True, failure_type="TaskFailed")
+        sf = _make_sf_client(status="FAILED", events=events)
+        with patch("idp_common.monitoring.stepfunctions_service.boto3") as mock_boto3:
+            mock_boto3.client.return_value = sf
+            timeline = analyze_execution_timeline(_EXEC_ARN)
+
+        # The last entered state before TaskFailed is ExtractDocument
+        assert timeline["failed_state"] == "ExtractDocument"
+        failed_entries = [s for s in timeline["states"] if s["is_failure"]]
+        assert len(failed_entries) == 1
+        assert failed_entries[0]["name"] == "ExtractDocument"
+        assert failed_entries[0]["name"] != ""  # must not be empty string
+
+    def test_task_timed_out_event_records_state_name(self):
+        """
+        TaskTimedOut does not carry stateExitedEventDetails either.
+        """
+        events = _make_events(include_failure=True, failure_type="TaskTimedOut")
+        sf = _make_sf_client(status="FAILED", events=events)
+        with patch("idp_common.monitoring.stepfunctions_service.boto3") as mock_boto3:
+            mock_boto3.client.return_value = sf
+            timeline = analyze_execution_timeline(_EXEC_ARN)
+
+        assert timeline["failed_state"] == "ExtractDocument"
+        failed_entries = [s for s in timeline["states"] if s["is_failure"]]
+        assert failed_entries[0]["name"] == "ExtractDocument"
+
 
 # ---------------------------------------------------------------------------
 # extract_failure_details — all 6 failure event types