lint

Author: nhulston

tests/test_api.py

@@ -7,17 +7,21 @@
 class TestDatadogLambdaAPI(unittest.TestCase):
     def setUp(self):
         api.api_key = None
-        self.env_patcher = patch.dict(os.environ, {
-            "DD_API_KEY_SECRET_ARN": "",
-            "DD_API_KEY_SSM_NAME": "",
-            "DD_KMS_API_KEY": "",
-            "DD_API_KEY": "",
-            "DATADOG_API_KEY": "",
-            "AWS_REGION": "",
-        }, clear=True)
+        self.env_patcher = patch.dict(
+            os.environ,
+            {
+                "DD_API_KEY_SECRET_ARN": "",
+                "DD_API_KEY_SSM_NAME": "",
+                "DD_KMS_API_KEY": "",
+                "DD_API_KEY": "",
+                "DATADOG_API_KEY": "",
+                "AWS_REGION": "",
+            },
+            clear=True,
+        )
         self.env_patcher.start()
 
-    @patch('boto3.client')
+    @patch("boto3.client")
     def test_secrets_manager_fips_endpoint(self, mock_boto3_client):
         mock_client = MagicMock()
         mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"}
@@ -30,14 +34,16 @@ def test_secrets_manager_fips_endpoint(self, mock_boto3_client):
 
         mock_boto3_client.assert_called_with(
             "secretsmanager",
-            endpoint_url="https://secretsmanager-fips.us-gov-east-1.amazonaws.com"
+            endpoint_url="https://secretsmanager-fips.us-gov-east-1.amazonaws.com",
         )
         self.assertEqual(api_key, "test-api-key")
 
-    @patch('boto3.client')
+    @patch("boto3.client")
     def test_ssm_fips_endpoint(self, mock_boto3_client):
         mock_client = MagicMock()
-        mock_client.get_parameter.return_value = {"Parameter": {"Value": "test-api-key"}}
+        mock_client.get_parameter.return_value = {
+            "Parameter": {"Value": "test-api-key"}
+        }
         mock_boto3_client.return_value = mock_client
 
         os.environ["AWS_REGION"] = "us-gov-west-1"
@@ -46,13 +52,12 @@ def test_ssm_fips_endpoint(self, mock_boto3_client):
         api_key = api.get_api_key()
 
         mock_boto3_client.assert_called_with(
-            "ssm",
-            endpoint_url="https://ssm-fips.us-gov-west-1.amazonaws.com"
+            "ssm", endpoint_url="https://ssm-fips.us-gov-west-1.amazonaws.com"
         )
         self.assertEqual(api_key, "test-api-key")
 
-    @patch('boto3.client')
-    @patch('datadog_lambda.api.decrypt_kms_api_key')
+    @patch("boto3.client")
+    @patch("datadog_lambda.api.decrypt_kms_api_key")
     def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client):
         mock_client = MagicMock()
         mock_boto3_client.return_value = mock_client
@@ -64,21 +69,12 @@ def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client):
         api_key = api.get_api_key()
 
         mock_boto3_client.assert_called_with(
-            "kms",
-            endpoint_url="https://kms-fips.us-gov-west-1.amazonaws.com"
+            "kms", endpoint_url="https://kms-fips.us-gov-west-1.amazonaws.com"
         )
         self.assertEqual(api_key, "test-api-key")
 
-    @patch('boto3.client')
+    @patch("boto3.client")
     def test_no_fips_for_standard_regions(self, mock_boto3_client):
         mock_client = MagicMock()
         mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"}
         mock_boto3_client.return_value = mock_client
-
-        os.environ.clear()
-        os.environ["AWS_REGION"] = "us-west-2"
-        os.environ["DD_API_KEY_SECRET_ARN"] = "test-arn"
-
-        api.get_api_key()
-
-        mock_boto3_client.assert_called_with("secretsmanager", endpoint_url=None)