From cfb3cced0c11f37d687eb9fa3dc6f19b6746ba6d Mon Sep 17 00:00:00 2001
From: Taegyun Kim <taegyun.kim@datadoghq.com>
Date: Fri, 27 Sep 2024 19:14:01 +0000
Subject: [PATCH] test for reversed locations

---
 tests/profiling/collector/pprof_utils.py   | 24 +++++++-
 tests/profiling_v2/collector/test_stack.py | 65 +++++++++++++++++++++-
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/tests/profiling/collector/pprof_utils.py b/tests/profiling/collector/pprof_utils.py
index ddf4dff74b..afee8c7cbd 100644
--- a/tests/profiling/collector/pprof_utils.py
+++ b/tests/profiling/collector/pprof_utils.py
@@ -28,6 +28,12 @@ def reinterpret_int_as_int64(value: int) -> int:
     return ctypes.c_int64(value).value
 
 
+class StackLocation:
+    def __init__(self, function_name: str, filename: str):
+        self.function_name = function_name
+        self.filename = filename
+
+
 class LockEventType(Enum):
     ACQUIRE = 1
     RELEASE = 2
@@ -50,7 +56,8 @@ def __init__(
 
 
 class StackEvent(EventBaseClass):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, locations: typing.Optional[typing.Any] = None, *args, **kwargs):
+        self.locations = locations
         super().__init__(*args, **kwargs)
 
 
@@ -266,4 +273,19 @@ def assert_lock_event(profile, sample: pprof_pb2.Sample, expected_event: LockEve
 
 
 def assert_stack_event(profile, sample: pprof_pb2.Sample, expected_event: StackEvent):
+    if expected_event.locations:
+        for location_id, expected_location in zip(sample.location_id, expected_event.locations):
+            location = get_location_with_id(profile, location_id)
+            function = get_function_with_id(profile, location.line[0].function_id)
+            assert (
+                profile.string_table[function.name] == expected_location.function_name
+            ), "Expected function {} got {}".format(
+                expected_location.function_name, profile.string_table[function.name]
+            )
+            assert profile.string_table[function.filename].endswith(
+                expected_location.filename
+            ), "Expected filepath to end with {}, got {}".format(
+                expected_location.filename, profile.string_table[function.filename]
+            )
+
     assert_base_event(profile, sample, expected_event)
diff --git a/tests/profiling_v2/collector/test_stack.py b/tests/profiling_v2/collector/test_stack.py
index cbd7c7c386..aeda1ddfb1 100644
--- a/tests/profiling_v2/collector/test_stack.py
+++ b/tests/profiling_v2/collector/test_stack.py
@@ -1,12 +1,75 @@
 import pytest
 
 
+@pytest.mark.subprocess(env=dict(DD_PROFILING_STACK_V2_ENABLED="true"))
+def test_stack_v2_locations():
+    import os
+    import time
+
+    from ddtrace.internal.datadog.profiling import ddup
+    from ddtrace.profiling.collector import stack
+    from tests.profiling.collector import pprof_utils
+
+    test_name = "test_locations"
+    pprof_prefix = "/tmp/" + test_name
+    output_filename = pprof_prefix + "." + str(os.getpid())
+
+    assert ddup.is_available
+    ddup.config(env="test", service=test_name, version="my_version", output_filename=pprof_prefix)
+    ddup.start()
+
+    def baz():
+        time.sleep(0.01)
+
+    def bar():
+        baz()
+
+    def foo():
+        bar()
+
+    with stack.StackCollector(None):
+        for _ in range(5):
+            foo()
+    ddup.upload()
+
+    profile = pprof_utils.parse_profile(output_filename)
+    samples = pprof_utils.get_samples_with_value_type(profile, "wall-time")
+    assert len(samples) > 0
+
+    locations_asserted = False
+    for sample in samples:
+        if len(sample.location_id) > 0:
+            pprof_utils.assert_stack_event(
+                profile,
+                sample,
+                expected_event=pprof_utils.StackEvent(
+                    locations=[
+                        pprof_utils.StackLocation(
+                            function_name="baz",
+                            filename="test_stack.py",
+                        ),
+                        pprof_utils.StackLocation(
+                            function_name="bar",
+                            filename="test_stack.py",
+                        ),
+                        pprof_utils.StackLocation(
+                            function_name="foo",
+                            filename="test_stack.py",
+                        ),
+                    ]
+                ),
+            )
+            locations_asserted = True
+
+    assert locations_asserted
+
+
 # Tests here are marked as subprocess as they are flaky when not marked as such,
 # similar to test_user_threads_have_native_id in test_threading.py. For some
 # reason, when the Span is created, it's not linked to the MainThread, and the
 # profiler can't find the corresponding Span for the MainThread.
 @pytest.mark.subprocess()
-def test_push_span():
+def test_push_span(span_type):
     import os
     import time
     import uuid