diff --git a/tests/web/rpc/test_aggregation.py b/tests/web/rpc/test_aggregation.py index 81fbb213a9..c7db14b9f9 100644 --- a/tests/web/rpc/test_aggregation.py +++ b/tests/web/rpc/test_aggregation.py @@ -1,14 +1,18 @@ +from typing import Any + import pytest from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ( AttributeAggregation, AttributeKey, ExtrapolationMode, Function, + Reliability, ) from snuba.web.rpc.common.aggregation import ( CUSTOM_COLUMN_PREFIX, CustomColumnInformation, + ExtrapolationContext, _get_closest_percentile_index, get_confidence_interval_column, ) @@ -77,6 +81,120 @@ def test_get_confidence_interval_column_for_non_extrapolatable_column() -> None: ) +@pytest.mark.parametrize( + ( + "row_data", + "column_name", + "average_sample_rate", + "reliability", + "extrapolated_data_present", + "is_extrapolated", + ), + [ + ( + { + "time": "2024-4-20 16:20:00", + "count(sentry.duration)": 100, + "p95(sentry.duration)": 123456, + CustomColumnInformation( + custom_column_id="confidence_interval", + referenced_column="count(sentry.duration)", + metadata={"function_type": "count"}, + ).to_alias(): 20, + CustomColumnInformation( + custom_column_id="average_sample_rate", + referenced_column="count(sentry.duration)", + metadata={}, + ).to_alias(): 0.5, + CustomColumnInformation( + custom_column_id="count", + referenced_column="count(sentry.duration)", + metadata={}, + ).to_alias(): 120, + "group_by_attr_1": "g1", + "group_by_attr_2": "a1", + }, + "count(sentry.duration)", + 0.5, + Reliability.RELIABILITY_HIGH, + True, + True, + ), + ( + { + "time": "2024-4-20 16:20:00", + "count(sentry.duration)": 100, + "min(sentry.duration)": 123456, + CustomColumnInformation( + custom_column_id="confidence_interval", + referenced_column="count(sentry.duration)", + metadata={"function_type": "count"}, + ).to_alias(): 20, + CustomColumnInformation( + custom_column_id="average_sample_rate", + referenced_column="count(sentry.duration)", + metadata={}, + ).to_alias(): 0.5, + CustomColumnInformation( + custom_column_id="count", + referenced_column="count(sentry.duration)", + metadata={}, + ).to_alias(): 120, + "group_by_attr_1": "g1", + "group_by_attr_2": "a1", + }, + "min(sentry.duration)", + 0, + Reliability.RELIABILITY_UNSPECIFIED, + False, + False, + ), + ( + { + "time": "2024-4-20 16:20:00", + "count(sentry.duration)": 100, + "p95(sentry.duration)": 123456, + CustomColumnInformation( + custom_column_id="confidence_interval", + referenced_column="count(sentry.duration)", + metadata={"function_type": "count"}, + ).to_alias(): 0, + CustomColumnInformation( + custom_column_id="average_sample_rate", + referenced_column="count(sentry.duration)", + metadata={}, + ).to_alias(): 0, + CustomColumnInformation( + custom_column_id="count", + referenced_column="count(sentry.duration)", + metadata={}, + ).to_alias(): 0, + "group_by_attr_1": "g1", + "group_by_attr_2": "a1", + }, + "count(sentry.duration)", + 0, + Reliability.RELIABILITY_UNSPECIFIED, + False, + True, + ), + ], +) +def test_get_extrapolation_meta( + row_data: dict[str, Any], + column_name: str, + average_sample_rate: float, + reliability: Reliability.ValueType, + extrapolated_data_present: bool, + is_extrapolated: bool, +) -> None: + extrapolation_context = ExtrapolationContext.from_row(column_name, row_data) + assert extrapolation_context.average_sample_rate == average_sample_rate + assert extrapolation_context.reliability == reliability + assert extrapolation_context.extrapolated_data_present == extrapolated_data_present + assert extrapolation_context.is_extrapolated == is_extrapolated + + @pytest.mark.parametrize( ("value", "percentile", "granularity", "width", "expected_index"), [