test(profiling): unflake a flaky test

tests/profiling/collector/pprof_utils.py

@@ -193,6 +193,25 @@ def get_function_with_id(profile: pprof_pb2.Profile, function_id: int) -> pprof_
     return next(function for function in profile.function if function.id == function_id)
 
 
+def get_location_from_id(profile: pprof_pb2.Profile, location_id: int) -> StackLocation:
+    """Get a StackLocation tuple from a location ID.
+
+    Args:
+        profile: The pprof profile containing location and function data
+        location_id: The location ID to look up
+
+    Returns:
+        A StackLocation with function_name, filename, and line_no
+    """
+    location = get_location_with_id(profile, location_id)
+    line = location.line[0]
+    function = get_function_with_id(profile, line.function_id)
+    function_name = profile.string_table[function.name]
+    filename = profile.string_table[function.filename]
+    line_no = line.line
+    return StackLocation(function_name=function_name, filename=filename, line_no=line_no)
+
+
 def assert_lock_events_of_type(
     profile: pprof_pb2.Profile,
     expected_events: List[LockEvent],
@@ -237,21 +256,25 @@ def assert_lock_events(
 
 
 def assert_str_label(string_table: Sequence[str], sample, key: str, expected_value: Optional[str]):
-    if expected_value:
-        label = get_label_with_key(string_table, sample, key)
-        # We use fullmatch to ensure that the whole string matches the expected value
-        # and not just a substring
-        assert label is not None, "Label {} not found in sample".format(key)
-        assert re.fullmatch(expected_value, string_table[label.str]), "Expected {} got {} for label {}".format(
-            expected_value, string_table[label.str], key
-        )
+    if not expected_value:
+        return
+
+    label = get_label_with_key(string_table, sample, key)
+    # We use fullmatch to ensure that the whole string matches the expected value
+    # and not just a substring
+    assert label is not None, f"Label {key} not found in sample"
+    assert re.fullmatch(expected_value, string_table[label.str]), (
+        f"Expected {expected_value} got {string_table[label.str]} for label {key}"
+    )
 
 
 def assert_num_label(string_table: Sequence[str], sample, key: str, expected_value: Optional[int]):
-    if expected_value:
-        label = get_label_with_key(string_table, sample, key)
-        assert label is not None, "Label {} not found in sample".format(key)
-        assert label.num == expected_value, "Expected {} got {} for label {}".format(expected_value, label.num, key)
+    if not expected_value:
+        return
+
+    label = get_label_with_key(string_table, sample, key)
+    assert label is not None, f"Label {key} not found in sample"
+    assert label.num == expected_value, f"Expected {expected_value} got {label.num} for label {key}"
 
 
 def assert_base_event(string_table: Sequence[str], sample: pprof_pb2.Sample, expected_event: EventBaseClass):
@@ -345,28 +368,38 @@ def assert_sample_has_locations(
         if DEBUG_TEST:
             print(loc)
 
-    assert checked, "Expected locations {} not found in sample locations: {}".format(
-        expected_locations, sample_loc_strs
-    )
+    assert checked, f"Expected locations {expected_locations} not found in sample locations: {sample_loc_strs}"
 
 
-def assert_stack_event(profile: pprof_pb2.Profile, sample: pprof_pb2.Sample, expected_event: StackEvent) -> None:
-    # Check that the sample has label "exception type" with value
-    assert_str_label(profile.string_table, sample, "exception type", expected_event.exception_type)
-    assert_sample_has_locations(profile, sample, expected_event.locations)
-    assert_base_event(profile.string_table, sample, expected_event)
+def assert_stack_event(
+    profile: pprof_pb2.Profile,
+    sample: pprof_pb2.Sample,
+    expected_event: StackEvent,
+    print_samples_on_failure: bool = False,
+) -> None:
+    try:
+        # Check that the sample has label "exception type" with value (no-op if expected_event.exception_type is None)
+        assert_str_label(profile.string_table, sample, "exception type", expected_event.exception_type)
+        assert_sample_has_locations(profile, sample, expected_event.locations)
+        assert_base_event(profile.string_table, sample, expected_event)
+    except AssertionError as e:
+        if print_samples_on_failure:
+            print_all_samples(profile)
+
+        raise e
 
 
 def assert_profile_has_sample(
     profile: pprof_pb2.Profile,
-    samples: List[pprof_pb2.Sample],
+    samples: Sequence[pprof_pb2.Sample],
     expected_sample: StackEvent,
     print_samples_on_failure: bool = False,
 ) -> None:
     found = False
     for sample in samples:
         try:
-            assert_stack_event(profile, sample, expected_sample)
+            # Pass print_samples_on_failure=False since we will print them later if needed
+            assert_stack_event(profile, sample, expected_sample, print_samples_on_failure=False)
             found = True
             break
         except AssertionError as e: