Replace module-level boto3 client with lazy initialization

Author: kris-szlapa

packages/syncKnowledgeBaseFunction/app/handler.py

@@ -5,7 +5,11 @@
 from app.config.config import KNOWLEDGEBASE_ID, DATA_SOURCE_ID
 
 logger = Logger(service="syncKnowledgeBaseFunction")
-bedrock_agent = boto3.client("bedrock-agent")
+
+
+def get_bedrock_agent():
+    """Get or create bedrock agent client"""
+    return boto3.client("bedrock-agent")
 
 
 @logger.inject_lambda_context
@@ -71,6 +75,7 @@ def handler(event, context):
                 )
 
                 ingestion_start_time = time.time()
+                bedrock_agent = get_bedrock_agent()
                 response = bedrock_agent.start_ingestion_job(
                     knowledgeBaseId=KNOWLEDGEBASE_ID,
                     dataSourceId=DATA_SOURCE_ID,

packages/syncKnowledgeBaseFunction/tests/test_handler.py

@@ -7,7 +7,7 @@
 @pytest.fixture
 def mock_env():
     """Mock environment variables"""
-    env_vars = {"KNOWLEDGEBASE_ID": "test-kb-id", "DATA_SOURCE_ID": "test-ds-id"}
+    env_vars = {"KNOWLEDGEBASE_ID": "test-kb-id", "DATA_SOURCE_ID": "test-ds-id", "AWS_REGION": "eu-west-2"}
     with patch.dict(os.environ, env_vars):
         yield env_vars
 
@@ -63,11 +63,12 @@ def multiple_s3_event():
     }
 
 
-@patch("app.handler.bedrock_agent")
+@patch("app.handler.get_bedrock_agent")
 @patch("time.time")
-def test_handler_success(mock_time, mock_bedrock, mock_env, lambda_context, s3_event):
+def test_handler_success(mock_time, mock_get_bedrock, mock_env, lambda_context, s3_event):
     """Test successful handler execution"""
     mock_time.side_effect = [1000, 1001, 1002, 1003]
+    mock_bedrock = mock_get_bedrock.return_value
     mock_bedrock.start_ingestion_job.return_value = {
         "ingestionJob": {"ingestionJobId": "job-123", "status": "STARTING"}
     }
@@ -85,11 +86,12 @@ def test_handler_success(mock_time, mock_bedrock, mock_env, lambda_context, s3_e
     )
 
 
-@patch("app.handler.bedrock_agent")
+@patch("app.handler.get_bedrock_agent")
 @patch("time.time")
-def test_handler_multiple_files(mock_time, mock_bedrock, mock_env, lambda_context, multiple_s3_event):
+def test_handler_multiple_files(mock_time, mock_get_bedrock, mock_env, lambda_context, multiple_s3_event):
     """Test handler with multiple S3 records"""
     mock_time.side_effect = [1000, 1001, 1002, 1003, 1004, 1005]
+    mock_bedrock = mock_get_bedrock.return_value
     mock_bedrock.start_ingestion_job.return_value = {
         "ingestionJob": {"ingestionJobId": "job-123", "status": "STARTING"}
     }
@@ -103,15 +105,16 @@ def test_handler_multiple_files(mock_time, mock_bedrock, mock_env, lambda_contex
     assert mock_bedrock.start_ingestion_job.call_count == 2
 
 
-@patch("app.handler.bedrock_agent")
+@patch("app.handler.get_bedrock_agent")
 @patch("time.time")
-def test_handler_conflict_exception(mock_time, mock_bedrock, mock_env, lambda_context, s3_event):
+def test_handler_conflict_exception(mock_time, mock_get_bedrock, mock_env, lambda_context, s3_event):
     """Test handler with ConflictException (job already running)"""
     mock_time.side_effect = [1000, 1001, 1002]
     error = ClientError(
         error_response={"Error": {"Code": "ConflictException", "Message": "Job already running"}},
         operation_name="StartIngestionJob",
     )
+    mock_bedrock = mock_get_bedrock.return_value
     mock_bedrock.start_ingestion_job.side_effect = error
 
     from app.handler import handler
@@ -122,15 +125,16 @@ def test_handler_conflict_exception(mock_time, mock_bedrock, mock_env, lambda_co
     assert "Files uploaded successfully - processing by existing ingestion job" in result["body"]
 
 
-@patch("app.handler.bedrock_agent")
+@patch("app.handler.get_bedrock_agent")
 @patch("time.time")
-def test_handler_aws_error(mock_time, mock_bedrock, mock_env, lambda_context, s3_event):
+def test_handler_aws_error(mock_time, mock_get_bedrock, mock_env, lambda_context, s3_event):
     """Test handler with other AWS error"""
     mock_time.side_effect = [1000, 1001, 1002]
     error = ClientError(
         error_response={"Error": {"Code": "AccessDenied", "Message": "Access denied"}},
         operation_name="StartIngestionJob",
     )
+    mock_bedrock = mock_get_bedrock.return_value
     mock_bedrock.start_ingestion_job.side_effect = error
 
     from app.handler import handler
@@ -141,11 +145,12 @@ def test_handler_aws_error(mock_time, mock_bedrock, mock_env, lambda_context, s3
     assert "AWS error: AccessDenied - Access denied" in result["body"]
 
 
-@patch("app.handler.bedrock_agent")
+@patch("app.handler.get_bedrock_agent")
 @patch("time.time")
-def test_handler_unexpected_error(mock_time, mock_bedrock, mock_env, lambda_context, s3_event):
+def test_handler_unexpected_error(mock_time, mock_get_bedrock, mock_env, lambda_context, s3_event):
     """Test handler with unexpected error"""
-    mock_time.side_effect = [1000, 1001, 1002]  # Added extra time call
+    mock_time.side_effect = [1000, 1001, 1002]
+    mock_bedrock = mock_get_bedrock.return_value
     mock_bedrock.start_ingestion_job.side_effect = Exception("Unexpected error")
 
     from app.handler import handler