add org tests back

getsentry · Jan 3, 2025 · ac0a0a4 · ac0a0a4
1 parent 6abb54e
commit ac0a0a4
Showing 1 changed file with 118 additions and 0 deletions.
diff --git 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"),
     [