Group magicmocks together

uktrade · Dec 13, 2024 · 993b0f8 · 993b0f8
1 parent d796718
commit 993b0f8
Showing 1 changed file with 5 additions and 9 deletions.
diff --git a/application-load-balancer/lambda_function/test_rotate_secret_lambda.py b/application-load-balancer/lambda_function/test_rotate_secret_lambda.py
@@ -44,16 +44,15 @@ def test_assumes_correct_role_for_cloudfront_access(self, rotator):
             }
         }
 
+        mock_sts = MagicMock()
+        mock_cloudfront = MagicMock()
+
+        mock_sts.assume_role.return_value = mock_credentials
+
         with patch('boto3.client') as mock_boto3_client:
-            # Set up separate mocks for STS and CloudFront clients
-            mock_sts = MagicMock()
-            mock_cloudfront = MagicMock()
-
             mock_boto3_client.side_effect = lambda service, **kwargs: \
                 mock_sts if service == 'sts' else mock_cloudfront
 
-            mock_sts.assume_role.return_value = mock_credentials
-
             client = rotator.get_cloudfront_client()
 
             mock_sts.assume_role.assert_called_once_with(
@@ -996,7 +995,6 @@ def test_handles_empty_distribution_list_gracefully(self, rotator):
         rotator.get_deployed_distributions = MagicMock()
         rotator.get_deployed_distributions.return_value = []
 
-        boto3.client = MagicMock()
         mock_boto_client = MagicMock()
         mock_boto_client.assume_role.return_value = mock_credentials
         mock_boto_client.get_secret_value.side_effect = [
@@ -1005,8 +1003,6 @@ def test_handles_empty_distribution_list_gracefully(self, rotator):
         ]
         mock_boto_client.describe_secret.return_value = mock_metadata
 
-        boto3.client.return_value = mock_boto_client
-
         rotator.get_waf_acl = MagicMock()
         rotator.update_cf_distro = MagicMock()
         time.sleep = MagicMock()