refactor: reuse the Lambda client

src/aws_durable_execution_sdk_python/lambda_service.py

@@ -937,19 +937,33 @@ def get_execution_state(
 class LambdaClient(DurableServiceClient):
     """Persist durable operations to the Lambda Durable Function APIs."""
 
+    boto3_instance: LambdaClient | None = None
+
     def __init__(self, client: Any) -> None:
         self.client = client
 
-    @staticmethod
-    def initialize_client() -> LambdaClient:
-        client = boto3.client(
-            "lambda",
-            config=Config(
-                connect_timeout=5,
-                read_timeout=50,
-            ),
-        )
-        return LambdaClient(client=client)
+    @classmethod
+    def initialize_client(cls) -> LambdaClient:
+        """Initialize or return cached Lambda client.
+
+        Implements lazy initialization with class-level caching to optimize
+        Lambda warm starts. The client is created once and reused across
+        invocations, avoiding repeated credential resolution and connection
+        pool setup.
+
+        Returns:
+            LambdaClient: The cached or newly created Lambda client.
+        """
+        if cls.boto3_instance is None:
+            client = boto3.client(
+                "lambda",
+                config=Config(
+                    connect_timeout=5,
+                    read_timeout=50,
+                ),
+            )
+            cls.boto3_instance = cls(client=client)
+        return cls.boto3_instance
 
     def checkpoint(
         self,

tests/lambda_service_test.py

@@ -37,6 +37,19 @@
     WaitOptions,
 )
 
+# =============================================================================
+# Fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def reset_lambda_client_cache():
+    """Reset the class-level Lambda client cache before and after each test."""
+    LambdaClient.boto3_instance = None
+    yield
+    LambdaClient.boto3_instance = None
+
+
 # =============================================================================
 # Tests for Data Classes (ExecutionDetails, ContextDetails, ErrorObject, etc.)
 # =============================================================================
@@ -1908,7 +1921,9 @@ def test_lambda_client_constructor():
 
 @patch.dict("os.environ", {}, clear=True)
 @patch("boto3.client")
-def test_lambda_client_initialize_client_default(mock_boto_client):
+def test_lambda_client_initialize_client_default(
+    mock_boto_client, reset_lambda_client_cache
+):
     """Test LambdaClient.initialize_client with default endpoint."""
     mock_client = Mock()
     mock_boto_client.return_value = mock_client
@@ -1928,7 +1943,9 @@ def test_lambda_client_initialize_client_default(mock_boto_client):
 
 @patch.dict("os.environ", {"AWS_ENDPOINT_URL_LAMBDA": "http://localhost:3000"})
 @patch("boto3.client")
-def test_lambda_client_initialize_client_with_endpoint(mock_boto_client):
+def test_lambda_client_initialize_client_with_endpoint(
+    mock_boto_client, reset_lambda_client_cache
+):
     """Test LambdaClient.initialize_client with custom endpoint (boto3 handles it automatically)."""
     mock_client = Mock()
     mock_boto_client.return_value = mock_client
@@ -2006,7 +2023,9 @@ def test_checkpoint_error_handling():
 
 @patch.dict("os.environ", {}, clear=True)
 @patch("boto3.client")
-def test_lambda_client_initialize_client_no_endpoint(mock_boto_client):
+def test_lambda_client_initialize_client_no_endpoint(
+    mock_boto_client, reset_lambda_client_cache
+):
     """Test LambdaClient.initialize_client without AWS_ENDPOINT_URL_LAMBDA."""
     mock_client = Mock()
     mock_boto_client.return_value = mock_client
@@ -2042,3 +2061,59 @@ def test_lambda_client_checkpoint_with_non_none_client_token():
     call_args = mock_client.checkpoint_durable_execution.call_args[1]
     assert call_args["ClientToken"] == "client_token_123"
     assert result.checkpoint_token == "new_token"  # noqa: S105
+
+
+# =============================================================================
+# Tests for LambdaClient caching behavior
+# =============================================================================
+
+
+@patch("boto3.client")
+def test_lambda_client_cache_reuses_client(mock_boto_client, reset_lambda_client_cache):
+    """Test that initialize_client returns the same cached client on subsequent calls."""
+    mock_client = Mock()
+    mock_boto_client.return_value = mock_client
+
+    # First call should create the client
+    client1 = LambdaClient.initialize_client()
+
+    # Second call should return the same cached client
+    client2 = LambdaClient.initialize_client()
+
+    # boto3.client should only be called once
+    mock_boto_client.assert_called_once()
+
+    # Both calls should return the same instance
+    assert client1 is client2
+
+
+@patch("boto3.client")
+def test_lambda_client_cache_creates_client_only_once(
+    mock_boto_client, reset_lambda_client_cache
+):
+    """Test that boto3.client is called only once even with multiple initialize_client calls."""
+    mock_client = Mock()
+    mock_boto_client.return_value = mock_client
+
+    # Call initialize_client multiple times
+    for _ in range(5):
+        LambdaClient.initialize_client()
+
+    # boto3.client should only be called once
+    assert mock_boto_client.call_count == 1
+
+
+@patch("boto3.client")
+def test_lambda_client_cache_is_class_level(
+    mock_boto_client, reset_lambda_client_cache
+):
+    """Test that the cache is stored at class level and persists across calls."""
+    mock_client = Mock()
+    mock_boto_client.return_value = mock_client
+
+    # Create client
+    client = LambdaClient.initialize_client()
+
+    # Verify the cache is set at class level
+    assert LambdaClient.boto3_instance is client
+    assert LambdaClient.boto3_instance is not None