Fix job activation and update tests

Author: aBurmeseDev

python/example_code/s3/scenarios/batch/s3_batch.py

@@ -1,6 +1,7 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+# snippet-start:[python.example_code.s3.S3Batch.scenario]
 """
 This module provides functionality for AWS S3 Batch Operations.
 It includes classes for managing CloudFormation stacks and S3 batch scenarios.
@@ -57,6 +58,9 @@ def deploy_cloudformation_stack(self, stack_name):
                                     }
                                 ]
                             },
+                            "ManagedPolicyArns": [
+                                "arn:aws:iam::aws:policy/AmazonS3FullAccess"
+                            ],
                             "Policies": [
                                 {
                                     "PolicyName": "S3BatchOperationsPolicy",
@@ -66,18 +70,8 @@ def deploy_cloudformation_stack(self, stack_name):
                                             {
                                                 "Effect": "Allow",
                                                 "Action": [
-                                                    "s3:GetObject",
-                                                    "s3:PutObject",
-                                                    "s3:PutObjectTagging",
-                                                    "s3:GetObjectTagging"
-                                                ],
-                                                "Resource": "*"
-                                            },
-                                            {
-                                                "Effect": "Allow",
-                                                "Action": [
-                                                    "s3:GetObject",
-                                                    "s3:PutObject"
+                                                    "s3:*",
+                                                    "s3-object-lambda:*"
                                                 ],
                                                 "Resource": "*"
                                             }
@@ -336,11 +330,12 @@ def create_s3_batch_job(self, account_id, role_arn, manifest_location,
                 Priority=10,
                 RoleArn=role_arn,
                 Description='Batch job for tagging objects',
-                ConfirmationRequired=True
+                # Set to False to avoid confirmation requirement
+                ConfirmationRequired=False
             )
             job_id = response['JobId']
             print(f"Created batch job with ID: {job_id}")
-            print("Job requires confirmation before starting...")
+            print("Job created and should start automatically")
             return job_id
         except ClientError as e:
             print(f"Error creating batch job: {e}")
@@ -403,6 +398,11 @@ def wait_for_job_ready(self, job_id, account_id, desired_status='Ready'):
                 print(f"Current job status: {current_status}")
                 if current_status == desired_status:
                     return True
+                # For jobs with ConfirmationRequired=True, they start in Suspended state
+                # and need to be activated
+                if current_status == 'Suspended':
+                    print("Job is in Suspended state, can proceed with activation")
+                    return True
                 if current_status in ['Active', 'Failed', 'Cancelled', 'Complete']:
                     print(f"Job is in {current_status} state, cannot update priority")
                     if 'FailureReasons' in response['Job']:
@@ -431,21 +431,40 @@ def update_job_priority(self, job_id, account_id):
             ClientError: If updating job priority fails
         """
         try:
-            if self.wait_for_job_ready(job_id, account_id):
+            # Check current job status
+            response = self.s3control_client.describe_job(
+                AccountId=account_id,
+                JobId=job_id
+            )
+            current_status = response['Job']['Status']
+            print(f"Current job status before update: {current_status}")
+            print(f"Full job details: {response['Job']}")
+            
+            # First try to update the job priority
+            try:
                 self.s3control_client.update_job_priority(
                     AccountId=account_id,
                     JobId=job_id,
                     Priority=20
                 )
-                print(f"Updated priority for job {job_id}")
+                print(f"Successfully updated priority for job {job_id}")
+            except ClientError as e:
+                print(f"Warning: Could not update job priority: {e}")
+                # Continue anyway to try activating the job
+            
+            # Then try to activate the job
+            try:
                 self.s3control_client.update_job_status(
                     AccountId=account_id,
                     JobId=job_id,
                     RequestedJobStatus='Active'
                 )
-                print("Job confirmed and started")
-            else:
-                print("Could not update job priority as job is not in Ready state")
+                print(f"Successfully activated job {job_id}")
+            except ClientError as e:
+                print(f"Error activating job: {e}")
+                if 'Message' in str(e):
+                    print(f"Detailed error message: {e.response.get('Message', '')}")
+                raise
         except ClientError as e:
             print(f"Error updating job priority: {e}")
             raise
@@ -579,8 +598,21 @@ def main():
                 raise ValueError("Job failed, stopping execution")
 
         wait_for_input()
-        print("\n2. Updating job priority...")
-        scenario.update_job_priority(job_id, account_id)
+        print("\n2. Checking job status...")
+        # Get current job status instead of trying to update priority
+        response = scenario.s3control_client.describe_job(
+            AccountId=account_id,
+            JobId=job_id
+        )
+        current_status = response['Job']['Status']
+        print(f"Current job status: {current_status}")
+        
+        # Only try to update priority if job is not already active
+        if current_status not in ['Active', 'Complete']:
+            print("\nUpdating job priority...")
+            scenario.update_job_priority(job_id, account_id)
+        else:
+            print("Job is already active or complete, no need to update priority.")
 
         print("\nCleanup")
         if input(
@@ -599,3 +631,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+# snippet-end:[python.example_code.s3.S3Batch.scenario]

python/example_code/s3/scenarios/batch/test/test_s3_batch.py

@@ -1,10 +1,12 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import pytest
-from unittest.mock import Mock, patch, MagicMock
-from botocore.exceptions import ClientError, WaiterError
+"""Unit tests for S3 batch operations module."""
+
 import json
+import pytest
+from unittest.mock import Mock, patch
+from botocore.exceptions import ClientError
 
 from s3_batch import CloudFormationHelper, S3BatchScenario, setup_resources
 
@@ -20,7 +22,7 @@ def cfn_helper(self):
     @patch('boto3.client')
     def test_init(self, mock_boto3_client):
         """Test CloudFormationHelper initialization."""
-        helper = CloudFormationHelper('us-east-1')
+        CloudFormationHelper('us-east-1')
         mock_boto3_client.assert_called_with('cloudformation', region_name='us-east-1')
 
     @patch('boto3.client')
@@ -29,14 +31,17 @@ def test_deploy_cloudformation_stack_success(self, mock_boto3_client, cfn_helper
         mock_client = Mock()
         mock_boto3_client.return_value = mock_client
         cfn_helper.cfn_client = mock_client
-        
         with patch.object(cfn_helper, '_wait_for_stack_completion'):
             cfn_helper.deploy_cloudformation_stack('test-stack')
-            
         mock_client.create_stack.assert_called_once()
         call_args = mock_client.create_stack.call_args
         assert call_args[1]['StackName'] == 'test-stack'
         assert 'CAPABILITY_IAM' in call_args[1]['Capabilities']
+        
+        # Verify the template includes AmazonS3FullAccess policy
+        template_body = json.loads(call_args[1]['TemplateBody'])
+        assert 'ManagedPolicyArns' in template_body['Resources']['S3BatchRole']['Properties']
+        assert 'arn:aws:iam::aws:policy/AmazonS3FullAccess' in template_body['Resources']['S3BatchRole']['Properties']['ManagedPolicyArns']
 
     @patch('boto3.client')
     def test_deploy_cloudformation_stack_failure(self, mock_boto3_client, cfn_helper):
@@ -48,7 +53,6 @@ def test_deploy_cloudformation_stack_failure(self, mock_boto3_client, cfn_helper
         )
         mock_boto3_client.return_value = mock_client
         cfn_helper.cfn_client = mock_client
-        
         with pytest.raises(ClientError):
             cfn_helper.deploy_cloudformation_stack('test-stack')
 
@@ -61,7 +65,7 @@ def test_get_stack_outputs_success(self, mock_boto3_client, cfn_helper):
                 'Outputs': [
                     {'OutputKey': 'S3BatchRoleArn', 'OutputValue': 'arn:aws:iam::123456789012:role/test-role'}
                 ]
-            }]
+        }]
         }
         mock_boto3_client.return_value = mock_client
         cfn_helper.cfn_client = mock_client
@@ -82,6 +86,7 @@ def test_destroy_cloudformation_stack_success(self, mock_boto3_client, cfn_helpe
         mock_client.delete_stack.assert_called_once_with(StackName='test-stack')
 
 
+
 class TestS3BatchScenario:
     """Test cases for S3BatchScenario class."""
 
@@ -164,6 +169,10 @@ def test_create_s3_batch_job_success(self, mock_boto3_client, s3_scenario):
 
         assert job_id == 'test-job-id'
         mock_s3control_client.create_job.assert_called_once()
+        
+        # Verify ConfirmationRequired is set to False
+        call_args = mock_s3control_client.create_job.call_args
+        assert call_args[1]['ConfirmationRequired'] is False
 
     @patch('boto3.client')
     def test_check_job_failure_reasons(self, mock_boto3_client, s3_scenario):
@@ -193,18 +202,68 @@ def test_wait_for_job_ready_success(self, mock_sleep, mock_boto3_client, s3_scen
         result = s3_scenario.wait_for_job_ready('test-job-id', '123456789012')
 
         assert result is True
+        
+    @patch('boto3.client')
+    @patch('time.sleep')
+    def test_wait_for_job_ready_suspended(self, mock_sleep, mock_boto3_client, s3_scenario):
+        """Test waiting for job with Suspended status."""
+        mock_s3control_client = Mock()
+        mock_s3control_client.describe_job.return_value = {
+            'Job': {'Status': 'Suspended'}
+        }
+        s3_scenario.s3control_client = mock_s3control_client
+        
+        result = s3_scenario.wait_for_job_ready('test-job-id', '123456789012')
+        
+        assert result is True
 
     @patch('boto3.client')
     def test_update_job_priority_success(self, mock_boto3_client, s3_scenario):
         """Test successful job priority update."""
         mock_s3control_client = Mock()
+        mock_s3control_client.describe_job.return_value = {
+            'Job': {'Status': 'Suspended'}
+        }
         s3_scenario.s3control_client = mock_s3control_client
 
-        with patch.object(s3_scenario, 'wait_for_job_ready', return_value=True):
-            s3_scenario.update_job_priority('test-job-id', '123456789012')
+        s3_scenario.update_job_priority('test-job-id', '123456789012')
 
         mock_s3control_client.update_job_priority.assert_called_once()
         mock_s3control_client.update_job_status.assert_called_once()
+        
+    @patch('boto3.client')
+    def test_update_job_priority_with_ready_status(self, mock_boto3_client, s3_scenario):
+        """Test job priority update with Ready status."""
+        mock_s3control_client = Mock()
+        mock_s3control_client.describe_job.return_value = {
+            'Job': {'Status': 'Ready'}
+        }
+        s3_scenario.s3control_client = mock_s3control_client
+        
+        s3_scenario.update_job_priority('test-job-id', '123456789012')
+            
+        mock_s3control_client.update_job_priority.assert_called_once()
+        mock_s3control_client.update_job_status.assert_called_once()
+        
+    @patch('boto3.client')
+    def test_update_job_priority_error_handling(self, mock_boto3_client, s3_scenario):
+        """Test error handling in job priority update."""
+        mock_s3control_client = Mock()
+        mock_s3control_client.describe_job.return_value = {
+            'Job': {'Status': 'Suspended'}
+        }
+        mock_s3control_client.update_job_priority.side_effect = ClientError(
+            {'Error': {'Code': 'InvalidRequest', 'Message': 'Cannot update priority'}},
+            'UpdateJobPriority'
+        )
+        mock_s3control_client.update_job_status = Mock()
+        s3_scenario.s3control_client = mock_s3control_client
+        
+        # Should not raise exception due to error handling
+        s3_scenario.update_job_priority('test-job-id', '123456789012')
+        
+        # Should still try to activate the job even if priority update fails
+        mock_s3control_client.update_job_status.assert_called_once()
 
     @patch('boto3.client')
     def test_cleanup_resources(self, mock_boto3_client, s3_scenario):
@@ -228,12 +287,14 @@ class TestUtilityFunctions:
     @patch('s3_batch.input', return_value='c')
     def test_wait_for_input_valid(self, mock_input):
         """Test wait_for_input with valid input."""
+        # pylint: disable=import-outside-toplevel
         from s3_batch import wait_for_input
         wait_for_input()  # Should not raise exception
 
     @patch('s3_batch.input', side_effect=['invalid', 'c'])
     def test_wait_for_input_invalid_then_valid(self, mock_input):
         """Test wait_for_input with invalid then valid input."""
+        # pylint: disable=import-outside-toplevel
         from s3_batch import wait_for_input
         wait_for_input()  # Should not raise exception