Replace LambdaContext class with a Proto and a dataclass

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,28 @@
 )
 
 
+@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
+
+
 # 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 (
+    LambdaContext,
+    StartDurableExecutionInput,
+)
 
 
 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"