chore(mql): Add test for metrics extrapolation (#76926)

Now that we are supporting extrapolation in generic metrics snuba, we
should a test in sentry to check this functionality

src/sentry/testutils/cases.py

@@ -1729,6 +1729,7 @@ class BaseMetricsTestCase(SnubaTestCase):
         timestamp: int,
         value: Any,
         aggregation_option: AggregationOption | None = None,
+        sampling_weight: int | None = None,
     ) -> None:
 
         parsed = parse_mri(mri)
@@ -1814,6 +1815,9 @@ class BaseMetricsTestCase(SnubaTestCase):
         if aggregation_option:
             msg["aggregation_option"] = aggregation_option.value
 
+        if sampling_weight:
+            msg["sampling_weight"] = sampling_weight
+
         if METRIC_PATH_MAPPING[use_case_id] == UseCaseKey.PERFORMANCE:
             entity = f"generic_metrics_{cls.ENTITY_SHORTHANDS[metric_type]}s"
         else:

tests/snuba/test_metrics_layer.py

@@ -73,16 +73,17 @@ class MQLTest(TestCase, BaseMetricsTestCase):
                 else:
                     value = i
                 self.store_metric(
-                    self.org_id,
-                    self.project.id,
-                    mri,
-                    {
+                    org_id=self.org_id,
+                    project_id=self.project.id,
+                    mri=mri,
+                    tags={
                         "transaction": f"transaction_{i % 2}",
                         "status_code": "500" if i % 3 == 0 else "200",
                         "device": "BlackBerry" if i % 2 == 0 else "Nokia",
                     },
-                    self.ts(self.hour_ago + timedelta(minutes=1 * i)),
-                    value,
+                    timestamp=self.ts(self.hour_ago + timedelta(minutes=1 * i)),
+                    value=value,
+                    sampling_weight=10,
                 )
         for mri, metric_type in self.metrics.items():
             assert metric_type in {"counter", "distribution", "set"}
@@ -873,6 +874,59 @@ class MQLTest(TestCase, BaseMetricsTestCase):
         for row in result["data"]:
             assert row["aggregate_value"] >= 86400
 
+    def test_extrapolated_generic_metrics(self) -> None:
+        query = MetricsQuery(
+            query=Timeseries(
+                metric=Metric(
+                    "transaction.duration",
+                    TransactionMRI.DURATION.value,
+                ),
+                aggregate="sum",
+                filters=[
+                    Condition(Column("status_code"), Op.EQ, "500"),
+                    Condition(Column("device"), Op.EQ, "BlackBerry"),
+                ],
+                groupby=[Column("transaction")],
+            ),
+            start=self.hour_ago,
+            end=self.now,
+            rollup=Rollup(interval=60, granularity=60),
+            scope=MetricsScope(
+                org_ids=[self.org_id],
+                project_ids=[self.project.id],
+                use_case_id=UseCaseID.TRANSACTIONS.value,
+            ),
+        )
+        request = Request(
+            dataset="generic_metrics",
+            app_id="tests",
+            query=query,
+            tenant_ids={"referrer": "metrics.testing.test", "organization_id": self.org_id},
+        )
+        result = run_query(request)
+        assert len(result["data"]) == 2
+        rows = result["data"]
+        assert rows[0]["aggregate_value"] in ([0], 0)
+        assert rows[0]["transaction"] == "transaction_0"
+        assert rows[1]["aggregate_value"] in ([6.00], 6)
+        assert rows[1]["transaction"] == "transaction_0"
+
+        # Set extrapolate flag to True. Since the sampling weight is set to 10, the extrapolated value should be 6*10
+        query = query.set_extrapolate(True)
+        request = Request(
+            dataset="generic_metrics",
+            app_id="tests",
+            query=query,
+            tenant_ids={"referrer": "metrics.testing.test", "organization_id": self.org_id},
+        )
+        result = run_query(request)
+        assert len(result["data"]) == 2
+        rows = result["data"]
+        assert rows[0]["aggregate_value"] in ([0], 0)
+        assert rows[0]["transaction"] == "transaction_0"
+        assert rows[1]["aggregate_value"] in ([60.00], 60)
+        assert rows[1]["transaction"] == "transaction_0"
+
 
 class MQLMetaTest(TestCase, BaseMetricsTestCase):
     def ts(self, dt: datetime) -> int: