Update partial transaction failure handling and add test (#1818)

elasticapm/contrib/serverless/aws.py

@@ -34,6 +34,7 @@
 import json
 import os
 import platform
+import re
 import time
 import urllib
 import warnings
@@ -453,7 +454,7 @@ def send_partial_transaction(self):
                     },
                 )
             except Exception as e:
-                if "HTTP 404" in str(e):
+                if re.match(r"HTTP [4,5]\d\d", str(e)):
                     REGISTER_PARTIAL_TRANSACTIONS = False
                     logger.info(
                         "APM Lambda Extension does not support partial transactions. "

tests/contrib/serverless/aws_tests.py

@@ -411,6 +411,9 @@ def test_func(event, context):
 
 
 def test_partial_transaction(event_api, context, sending_elasticapm_client):
+    import elasticapm.contrib.serverless.aws
+
+    elasticapm.contrib.serverless.aws.REGISTER_PARTIAL_TRANSACTIONS = True
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
     os.environ["ELASTIC_APM_LAMBDA_APM_SERVER"] = "http://localhost:8200"
 
@@ -427,3 +430,28 @@ def test_func(event, context):
     assert b"metadata" in request.data
     assert b"transaction" in request.data
     sending_elasticapm_client.close()
+
+
+def test_partial_transaction_failure(event_api, context, sending_elasticapm_client):
+    import elasticapm.contrib.serverless.aws
+
+    elasticapm.contrib.serverless.aws.REGISTER_PARTIAL_TRANSACTIONS = True
+    os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
+    os.environ["ELASTIC_APM_LAMBDA_APM_SERVER"] = "http://localhost:8200"
+    sending_elasticapm_client.httpserver.code = 404
+    sending_elasticapm_client.httpserver.content = "go away"
+
+    @capture_serverless
+    def test_func(event, context):
+        return {"statusCode": 200, "headers": {"foo": "bar"}}
+
+    test_func(event_api, context)
+    test_func(event_api, context)
+
+    assert len(sending_elasticapm_client.httpserver.requests) == 3
+    request = sending_elasticapm_client.httpserver.requests[0]
+    assert request.full_path == "/register/transaction?"
+    assert request.content_type == "application/vnd.elastic.apm.transaction+ndjson"
+    assert b"metadata" in request.data
+    assert b"transaction" in request.data
+    sending_elasticapm_client.close()