test(metrics): Add simple test for metrics aggregator flush

relay-metrics/src/aggregator.rs

@@ -897,6 +897,7 @@ fn validate_metric_tags(mut key: BucketKey, aggregator_config: &AggregatorConfig
 
 #[cfg(test)]
 mod tests {
+    use insta::assert_debug_snapshot;
     use similar_asserts::assert_eq;
 
     use super::*;
@@ -1249,6 +1250,84 @@ mod tests {
         assert_eq!(aggregator.cost_tracker.total_cost, 0);
     }
 
+    #[tokio::test]
+    async fn test_aggregator_flush() {
+        // Make sure that the right cost is added / subtracted
+        let mut aggregator: Aggregator = Aggregator::new(AggregatorConfig {
+            bucket_interval: 10,
+            ..test_config()
+        });
+        let project_key = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fed").unwrap();
+
+        let now = UnixTimestamp::now();
+        tokio::time::pause();
+
+        for i in 0..3u32 {
+            for (name, offset) in [("foo", 30), ("bar", 15)] {
+                let timestamp = now + Duration::from_secs(60 * u64::from(i) + offset);
+                let bucket = Bucket {
+                    timestamp,
+                    width: 0,
+                    name: format!("c:transactions/{name}_{i}@none").into(),
+                    value: BucketValue::counter(i.into()),
+                    tags: BTreeMap::new(),
+                    metadata: BucketMetadata::new(timestamp),
+                };
+
+                aggregator.merge(project_key, bucket, None).unwrap();
+            }
+        }
+
+        let mut flush_buckets = || {
+            let mut result = Vec::new();
+            for (partition, v) in aggregator.pop_flush_buckets(false) {
+                assert!(partition.is_none());
+                for (pk, buckets) in v {
+                    assert_eq!(pk, project_key);
+                    for bucket in buckets {
+                        result.push(bucket.name.to_string());
+                    }
+                }
+            }
+            result.sort();
+            result
+        };
+
+        assert!(flush_buckets().is_empty());
+
+        tokio::time::advance(Duration::from_secs(60)).await;
+        assert_debug_snapshot!(flush_buckets(), @r###"
+        [
+            "c:transactions/bar_0@none",
+            "c:transactions/foo_0@none",
+        ]
+        "###);
+        assert!(flush_buckets().is_empty());
+
+        tokio::time::advance(Duration::from_secs(60)).await;
+        assert_debug_snapshot!(flush_buckets(), @r###"
+        [
+            "c:transactions/bar_1@none",
+            "c:transactions/foo_1@none",
+        ]
+        "###);
+        assert!(flush_buckets().is_empty());
+
+        tokio::time::advance(Duration::from_secs(60)).await;
+        assert_debug_snapshot!(flush_buckets(), @r###"
+        [
+            "c:transactions/bar_2@none",
+            "c:transactions/foo_2@none",
+        ]
+        "###);
+        assert!(flush_buckets().is_empty());
+
+        tokio::time::advance(Duration::from_secs(3600)).await;
+        assert!(flush_buckets().is_empty());
+
+        assert!(aggregator.into_buckets().is_empty());
+    }
+
     #[test]
     fn test_get_bucket_timestamp_overflow() {
         let config = AggregatorConfig {