Add a Lambda Context Dataclass

Author: FullyTyped

src/aws_durable_execution_sdk_python_testing/invoker.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-import time
 from typing import TYPE_CHECKING, Any, Protocol
 
 import boto3  # type: ignore
@@ -11,11 +10,11 @@
     DurableExecutionInvocationOutput,
     InitialExecutionState,
 )
-from aws_durable_execution_sdk_python.lambda_context import LambdaContext
 
 from aws_durable_execution_sdk_python_testing.exceptions import (
     DurableFunctionsTestError,
 )
+from aws_durable_execution_sdk_python_testing.model import LambdaContext
 
 
 if TYPE_CHECKING:
@@ -46,12 +45,9 @@ def create_test_lambda_context() -> LambdaContext:
     }
 
     return LambdaContext(
-        invoke_id="test-invoke-12345",
+        aws_request_id="test-invoke-12345",
         client_context=client_context_dict,
-        cognito_identity=cognito_identity_dict,
-        epoch_deadline_time_in_ms=int(
-            (time.time() + 900) * 1000
-        ),  # 15 minutes from now
+        identity=cognito_identity_dict,
         invoked_function_arn="arn:aws:lambda:us-west-2:123456789012:function:test-function",
         tenant_id="test-tenant-789",
     )

src/aws_durable_execution_sdk_python_testing/model.py

@@ -25,6 +25,48 @@
 )
 
 
+@dataclass(frozen=True)
+class LambdaContext:
+    """Lambda context for testing."""
+
+    aws_request_id: str
+    log_group_name: str | None = None
+    log_stream_name: str | None = None
+    function_name: str | None = None
+    memory_limit_in_mb: str | None = None
+    function_version: str | None = None
+    invoked_function_arn: str | None = None
+    tenant_id: str | None = None
+    client_context: dict | None = None
+    identity: dict | None = None
+
+    def get_remaining_time_in_millis(self) -> int:
+        return 900000  # 15 minutes default
+
+    def log(self, msg) -> None:
+        pass  # No-op for testing
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        required_fields = ["aws_request_id"]
+        for field in required_fields:
+            if field not in data:
+                msg: str = f"Missing required field: {field}"
+                raise InvalidParameterValueException(msg)
+        return cls(
+            aws_request_id=data["aws_request_id"],
+            log_group_name=data.get("log_group_name"),
+            log_stream_name=data.get("log_stream_name"),
+            function_name=data.get("function_name"),
+            memory_limit_in_mb=data.get("memory_limit_in_mb"),
+            function_version=data.get("function_version"),
+            invoked_function_arn=data.get("invoked_function_arn"),
+            tenant_id=data.get("tenant_id"),
+            client_context=data.get("client_context"),
+            identity=data.get("identity"),
+        )
+
+
 # Web API specific models (not in Smithy but needed for web interface)
 @dataclass(frozen=True)
 class StartDurableExecutionInput:

tests/invoker_test.py

@@ -11,22 +11,23 @@
     InitialExecutionState,
     InvocationStatus,
 )
-from aws_durable_execution_sdk_python.lambda_context import LambdaContext
 
 from aws_durable_execution_sdk_python_testing.execution import Execution
 from aws_durable_execution_sdk_python_testing.invoker import (
     InProcessInvoker,
     LambdaInvoker,
     create_test_lambda_context,
 )
-from aws_durable_execution_sdk_python_testing.model import StartDurableExecutionInput
+from aws_durable_execution_sdk_python_testing.model import (
+    StartDurableExecutionInput,
+    LambdaContext,
+)
 
 
 def test_create_test_lambda_context():
     """Test creating a test lambda context."""
     context = create_test_lambda_context()
 
-    assert isinstance(context, LambdaContext)
     assert (
         context.invoked_function_arn
         == "arn:aws:lambda:us-west-2:123456789012:function:test-function"