|
| 1 | +from unittest import mock |
| 2 | + |
| 3 | +import pytest |
| 4 | + |
| 5 | +from posthog.tracing import _limits as limits_module |
| 6 | +from posthog.tracing._limits import ( |
| 7 | + bound_attributes, |
| 8 | + truncate_attribute_value, |
| 9 | + truncate_attributes, |
| 10 | +) |
| 11 | +from posthog.tracing._otlp import ( |
| 12 | + CIRCULAR_VALUE, |
| 13 | + MAX_VALUE_ITEMS, |
| 14 | + MAX_VALUE_NODES, |
| 15 | + TRUNCATED_VALUE, |
| 16 | + to_any_value, |
| 17 | +) |
| 18 | +from posthog.tracing._sanitize import FUNCTION_VALUE, UNSERIALIZABLE_VALUE |
| 19 | + |
| 20 | + |
| 21 | +class TestTruncateAttributeValue: |
| 22 | + def test_truncates_a_long_string(self): |
| 23 | + assert truncate_attribute_value("x" * 40000, 8192) == "x" * 8192 |
| 24 | + |
| 25 | + def test_never_walks_a_scalar(self): |
| 26 | + with mock.patch.object(limits_module, "_truncate") as walk: |
| 27 | + assert truncate_attribute_value("short", 8) == "short" |
| 28 | + assert truncate_attribute_value(7, 8) == 7 |
| 29 | + assert truncate_attribute_value(None, 8) is None |
| 30 | + assert not walk.called |
| 31 | + |
| 32 | + def test_returns_a_short_string_unchanged(self): |
| 33 | + assert truncate_attribute_value("short", 8192) == "short" |
| 34 | + |
| 35 | + @pytest.mark.parametrize("value", [42, 1.5, True, None, 2**70]) |
| 36 | + def test_leaves_numbers_booleans_and_none_alone(self, value): |
| 37 | + assert truncate_attribute_value(value, 3) == value |
| 38 | + |
| 39 | + def test_reaches_strings_nested_in_mappings_and_lists(self): |
| 40 | + value = {"body": "x" * 40000, "items": ["y" * 20, {"deep": "z" * 20}]} |
| 41 | + assert truncate_attribute_value(value, 8) == { |
| 42 | + "body": "x" * 8, |
| 43 | + "items": ["y" * 8, {"deep": "z" * 8}], |
| 44 | + } |
| 45 | + |
| 46 | + def test_does_not_mutate_the_callers_value(self): |
| 47 | + value = {"body": "x" * 20} |
| 48 | + truncate_attribute_value(value, 4) |
| 49 | + assert value == {"body": "x" * 20} |
| 50 | + |
| 51 | + def test_a_self_referencing_value_terminates_with_the_encoders_marker(self): |
| 52 | + value: dict = {"name": "n" * 20} |
| 53 | + value["self"] = value |
| 54 | + assert truncate_attribute_value(value, 4) == { |
| 55 | + "name": "nnnn", |
| 56 | + "self": CIRCULAR_VALUE, |
| 57 | + } |
| 58 | + |
| 59 | + def test_siblings_sharing_one_object_are_not_a_cycle(self): |
| 60 | + shared = {"k": "v" * 10} |
| 61 | + assert truncate_attribute_value([shared, shared], 2) == [ |
| 62 | + {"k": "vv"}, |
| 63 | + {"k": "vv"}, |
| 64 | + ] |
| 65 | + |
| 66 | + def test_marks_items_past_the_encoders_item_cap(self): |
| 67 | + bounded = truncate_attribute_value(["a"] * (MAX_VALUE_ITEMS + 5), 8) |
| 68 | + assert len(bounded) == MAX_VALUE_ITEMS + 1 |
| 69 | + assert bounded[-1] == TRUNCATED_VALUE |
| 70 | + # The encoder emits the same shape it would have for the original. |
| 71 | + assert to_any_value(bounded) == to_any_value(["a"] * (MAX_VALUE_ITEMS + 5)) |
| 72 | + |
| 73 | + def test_stringifies_and_bounds_a_type_the_encoder_would_stringify(self): |
| 74 | + class Big: |
| 75 | + def __str__(self): |
| 76 | + return "b" * 100 |
| 77 | + |
| 78 | + assert truncate_attribute_value(Big(), 10) == "b" * 10 |
| 79 | + assert truncate_attribute_value(b"\x00" * 100, 10) == "b'\\x00\\x00" |
| 80 | + |
| 81 | + def test_a_key_the_encoder_skips_does_not_spend_the_walks_budget(self): |
| 82 | + # The encoder drops "" without charging for its value, so the walk must |
| 83 | + # too, or "x" would ship unbounded once the walk's budget ran out. |
| 84 | + value = {"": list(range(999)), "a": [list(range(999))] * 9, "x": "A" * 1000} |
| 85 | + bounded = truncate_attribute_value(value, 100) |
| 86 | + assert "" not in bounded |
| 87 | + assert bounded["x"] == "A" * 100 |
| 88 | + encoded = to_any_value(bounded)["kvlistValue"]["values"] |
| 89 | + x = next(kv for kv in encoded if kv["key"] == "x") |
| 90 | + assert len(x["value"]["stringValue"]) == 100 |
| 91 | + |
| 92 | + def test_leaves_a_callable_for_the_encoders_marker(self): |
| 93 | + def handler(): |
| 94 | + pass |
| 95 | + |
| 96 | + assert truncate_attribute_value(handler, 3) is handler |
| 97 | + assert truncate_attribute_value({"fn": handler}, 3) == {"fn": handler} |
| 98 | + assert to_any_value(truncate_attribute_value(handler, 3)) == { |
| 99 | + "stringValue": FUNCTION_VALUE |
| 100 | + } |
| 101 | + |
| 102 | + def test_a_raising_str_costs_only_that_value(self): |
| 103 | + class Hostile: |
| 104 | + def __str__(self): |
| 105 | + raise RuntimeError("no") |
| 106 | + |
| 107 | + assert truncate_attribute_value({"a": Hostile(), "b": "ok"}, 8) == { |
| 108 | + "a": UNSERIALIZABLE_VALUE, |
| 109 | + "b": "ok", |
| 110 | + } |
| 111 | + |
| 112 | + def test_a_raising_accessor_costs_only_that_key(self): |
| 113 | + class Explosive(dict): |
| 114 | + def __getitem__(self, key): |
| 115 | + if key == "bad": |
| 116 | + raise RuntimeError("no") |
| 117 | + return super().__getitem__(key) |
| 118 | + |
| 119 | + assert truncate_attribute_value(Explosive(good="g" * 9, bad=1), 3) == { |
| 120 | + "good": "ggg", |
| 121 | + "bad": UNSERIALIZABLE_VALUE, |
| 122 | + } |
| 123 | + |
| 124 | + |
| 125 | +class TestBoundAttributes: |
| 126 | + def test_keeps_the_earliest_entries_and_counts_the_rest(self): |
| 127 | + source = {f"k{i}": i for i in range(130)} |
| 128 | + attributes, dropped = bound_attributes(source, 128, 8192) |
| 129 | + assert list(attributes) == [f"k{i}" for i in range(128)] |
| 130 | + assert dropped == 2 |
| 131 | + |
| 132 | + @pytest.mark.parametrize( |
| 133 | + "source", |
| 134 | + [{"a": None, "b": 1, "c": 2}, {"b": 1, "c": 2, "a": None}], |
| 135 | + ids=["before-the-cap", "past-the-cap"], |
| 136 | + ) |
| 137 | + def test_a_none_value_spends_no_slot_and_counts_no_drop(self, source): |
| 138 | + attributes, dropped = bound_attributes(source, 2, 8) |
| 139 | + assert attributes == {"b": 1, "c": 2} |
| 140 | + assert dropped == 0 |
| 141 | + |
| 142 | + def test_a_real_value_past_the_cap_counts_a_drop(self): |
| 143 | + attributes, dropped = bound_attributes({"b": 1, "c": 2, "a": 3}, 2, 8) |
| 144 | + assert attributes == {"b": 1, "c": 2} |
| 145 | + assert dropped == 1 |
| 146 | + |
| 147 | + def test_bounds_each_value(self): |
| 148 | + attributes, _ = bound_attributes({"a": "x" * 20}, 2, 5) |
| 149 | + assert attributes == {"a": "xxxxx"} |
| 150 | + |
| 151 | + def test_an_empty_key_spends_no_slot_and_counts_no_drop(self): |
| 152 | + attributes, dropped = bound_attributes({"": 1, "b": 2, "c": 3}, 2, 8) |
| 153 | + assert attributes == {"b": 2, "c": 3} |
| 154 | + assert dropped == 0 |
| 155 | + |
| 156 | + |
| 157 | +class TestTruncateAttributes: |
| 158 | + def test_bounds_every_value_as_a_copy(self): |
| 159 | + source = {"service.name": "api", "blob": "x" * 20} |
| 160 | + assert truncate_attributes(source, 4) == {"service.name": "api", "blob": "xxxx"} |
| 161 | + assert source["blob"] == "x" * 20 |
| 162 | + |
| 163 | + |
| 164 | +class TestWalkBounds: |
| 165 | + def test_walks_no_more_strings_than_the_encoder_would_emit(self): |
| 166 | + # A thousand paths to one shared list of a thousand strings. Leaves are |
| 167 | + # charged against the node budget, as in the encoder, so the walk does |
| 168 | + # not copy every string on every path. |
| 169 | + inner = ["x" * 50] * 1000 |
| 170 | + bounded = truncate_attribute_value([inner] * 1000, 8) |
| 171 | + walked = sum( |
| 172 | + 1 |
| 173 | + for items in bounded |
| 174 | + if items is not inner |
| 175 | + for item in items |
| 176 | + if item == "x" * 8 |
| 177 | + ) |
| 178 | + assert 0 < walked <= MAX_VALUE_NODES |
| 179 | + |
| 180 | + def test_stops_walking_a_mapping_at_the_encoders_item_cap(self): |
| 181 | + value = {f"k{i}": "v" * 50 for i in range(MAX_VALUE_ITEMS + 50)} |
| 182 | + bounded = truncate_attribute_value(value, 4) |
| 183 | + assert len(bounded) == MAX_VALUE_ITEMS |
0 commit comments