diff --git a/src/aws_durable_execution_sdk_python_testing/invoker.py b/src/aws_durable_execution_sdk_python_testing/invoker.py index b38f249..dfde61b 100644 --- a/src/aws_durable_execution_sdk_python_testing/invoker.py +++ b/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", ) diff --git a/src/aws_durable_execution_sdk_python_testing/model.py b/src/aws_durable_execution_sdk_python_testing/model.py index af24767..f13e47c 100644 --- a/src/aws_durable_execution_sdk_python_testing/model.py +++ b/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: diff --git a/tests/invoker_test.py b/tests/invoker_test.py index 9074496..8a40973 100644 --- a/tests/invoker_test.py +++ b/tests/invoker_test.py @@ -11,7 +11,6 @@ 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 ( @@ -19,14 +18,16 @@ 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"