Add Serverless Framework Support Revision 3 (#127)

* Mimic Segments now perform no-op on segment-only methods.
* Added comments to address this contract.
* Unit tests to ensure segment-only methods don't add data
to the mimic segment. Serialization methods are also tested
to ensure this.

Author: chanchiem

aws_xray_sdk/core/models/mimic_segment.py

@@ -18,6 +18,9 @@ class MimicSegment(Segment):
     as a node on the service graph. For all purposes, the MimicSegment can be interacted as if it's
     a real segment, meaning that all methods that exist only in a Segment but not a subsegment
     is available to be used.
+
+    The following methods are no-ops and will not be sent to the service:
+    set_rule_name, set_service, and set_user
     """
 
     def __init__(self, facade_segment, original_segment):
@@ -28,10 +31,19 @@ def __init__(self, facade_segment, original_segment):
                                            traceid=facade_segment.trace_id, parent_id=facade_segment.id,
                                            sampled=facade_segment.sampled)
 
+    def set_rule_name(self, rule_name):
+        pass
+
+    def set_service(self, service_info):
+        pass
+
+    def set_user(self, user):
+        pass
+
     def __getstate__(self):
         """
-        Used during serialization. We mark the subsegment properties to let the dataplane know
-        that we want the mimic segment to be transformed as a subsegment.
+        Used during serialization. We mark the mimic segment as a subsegment to
+        let the X-Ray service know to create this mimic segment as a subsegment.
         """
         properties = super(MimicSegment, self).__getstate__()
         properties['type'] = 'subsegment'

aws_xray_sdk/core/serverless_lambda_context.py

@@ -18,7 +18,7 @@ class ServerlessLambdaContext(LambdaContext):
     creates a Segment masked as a Subsegment known as a MimicSegment underneath
     the Lambda-generated Facade Segment. This ensures that middleware->recorder's
     consequent calls to "put_segment()" will not throw exceptions but instead create
-    subsegments underneath the lambda-generated segment. This context also
+    subsegments underneath the lambda-generated Facade Segment. This context also
     ensures that FacadeSegments exist through underlying calls to _refresh_context().
     """
     def __init__(self, context_missing='RUNTIME_ERROR'):
@@ -31,7 +31,6 @@ def put_segment(self, segment):
         """
         Convert the segment into a mimic segment and append it to FacadeSegment's subsegment list.
         :param Segment segment:
-        :return:
         """
         # When putting a segment, convert it to a mimic segment and make it a child of the Facade Segment.
         parent_facade_segment = self.__get_facade_entity()  # type: FacadeSegment
@@ -102,8 +101,9 @@ def get_trace_entity(self):
 
     def set_trace_entity(self, trace_entity):
         """
-        Store the input trace_entity to local context. It will overwrite all
-        existing ones if there is any.
+        Stores the input trace_entity to local context. It will overwrite all
+        existing ones if there is any. If the entity passed in is a segment,
+        it will automatically be converted to a mimic segment.
         """
         if type(trace_entity) == Segment:
             # Convert to a mimic segment.

tests/test_mimic_segment.py

@@ -71,6 +71,8 @@ def test_facade_segment_properties():
 
 def test_segment_methods_on_mimic():
     # Test to make sure that segment methods exist and function for the Mimic Segment
+    # And ensure that the methods (other than get/set origin_trace_header) don't modify
+    # the segment.
     mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment)  # type: MimicSegment
     assert not getattr(mimic_segment, "service", None)
     assert not getattr(mimic_segment, "user", None)
@@ -82,20 +84,44 @@ def test_segment_methods_on_mimic():
     assert getattr(original_segment, "ref_counter", None)
     assert getattr(original_segment, "_subsegments_counter", None)
 
-    mimic_segment.set_service("SomeService")
-    original_segment.set_service("SomeService")
-    assert original_segment.service == original_segment.service
-
     assert original_segment.get_origin_trace_header() == mimic_segment.get_origin_trace_header()
     mimic_segment.save_origin_trace_header("someheader")
     original_segment.save_origin_trace_header("someheader")
     assert original_segment.get_origin_trace_header() == mimic_segment.get_origin_trace_header()
 
     # No exception is thrown
     test_dict = {"akey": "avalue"}
-    original_segment.set_aws(test_dict)
-    original_segment.set_rule_name(test_dict)
+    test_rule_name = {"arule": "name"}
+    original_segment.set_aws(test_dict.copy())
+    original_segment.set_rule_name(test_rule_name.copy())
     original_segment.set_user("SomeUser")
-    mimic_segment.set_aws(test_dict)
-    mimic_segment.set_rule_name(test_dict)
+    original_segment.set_service("SomeService")
+    mimic_segment.set_aws(test_dict.copy())
+    mimic_segment.set_rule_name(test_rule_name.copy())
     mimic_segment.set_user("SomeUser")
+    mimic_segment.set_service("SomeService")
+
+    # Original segment should contain these properties
+    # but not the mimic segment.
+    assert getattr(original_segment, "service", None)
+    assert getattr(original_segment, "user", None)
+    assert 'xray' in getattr(original_segment, "aws", None)
+    assert 'sampling_rule_name' in getattr(original_segment, "aws", None)['xray']
+
+    assert not getattr(mimic_segment, "service", None)
+    assert not getattr(mimic_segment, "user", None)
+    # Originally set by rule_name, but no-op so nothing is set.
+    assert 'xray' not in getattr(mimic_segment, "aws", None)
+
+    # Ensure serialization also doesn't have those properties in mimic but do in original
+    original_segment_serialized = original_segment.__getstate__()
+    mimic_segment_serialized = mimic_segment.__getstate__()
+
+    assert 'service' in original_segment_serialized
+    assert 'user' in original_segment_serialized
+    assert 'xray' in original_segment_serialized['aws']
+    assert 'sampling_rule_name' in original_segment_serialized['aws']['xray']
+
+    assert 'service' not in mimic_segment_serialized
+    assert 'user' not in mimic_segment_serialized
+    assert 'xray' not in mimic_segment_serialized['aws']