From 2bd18b7fa53615f2b2d6a31f44c8978f7db3f1de Mon Sep 17 00:00:00 2001
From: Joris Bayer <joris.bayer@sentry.io>
Date: Thu, 19 Sep 2024 13:35:04 +0200
Subject: [PATCH] tests

---
 .../services/buffer/envelope_buffer/mod.rs    | 109 +++++-------------
 1 file changed, 32 insertions(+), 77 deletions(-)

diff --git a/relay-server/src/services/buffer/envelope_buffer/mod.rs b/relay-server/src/services/buffer/envelope_buffer/mod.rs
index 9ae469f2f6..d66841ebb6 100644
--- a/relay-server/src/services/buffer/envelope_buffer/mod.rs
+++ b/relay-server/src/services/buffer/envelope_buffer/mod.rs
@@ -697,6 +697,17 @@ mod tests {
         MemoryChecker::new(MemoryStat::default(), mock_config("my/db/path").clone())
     }
 
+    async fn peek_project_key(buffer: &mut EnvelopeBuffer<MemoryStackProvider>) -> ProjectKey {
+        buffer
+            .peek()
+            .await
+            .unwrap()
+            .envelope()
+            .unwrap()
+            .meta()
+            .public_key()
+    }
+
     #[tokio::test]
     async fn test_insert_pop() {
         let mut buffer = EnvelopeBuffer::<MemoryStackProvider>::new(mock_memory_checker());
@@ -712,110 +723,48 @@ mod tests {
             .push(new_envelope(project_key1, None, None))
             .await
             .unwrap();
-        assert_eq!(
-            buffer
-                .peek()
-                .await
-                .unwrap()
-                .envelope()
-                .unwrap()
-                .meta()
-                .public_key(),
-            project_key1
-        );
 
         buffer
             .push(new_envelope(project_key2, None, None))
             .await
             .unwrap();
+
+        // Both projects are ready, so project 2 is on top (has the newest envelopes):
+        assert_eq!(peek_project_key(&mut buffer).await, project_key2);
+
+        buffer.mark_ready(&project_key1, false);
+        buffer.mark_ready(&project_key2, false);
+
         // Both projects are not ready, so project 1 is on top (has the oldest envelopes):
-        assert_eq!(
-            buffer
-                .peek()
-                .await
-                .unwrap()
-                .envelope()
-                .unwrap()
-                .meta()
-                .public_key(),
-            project_key1
-        );
+        assert_eq!(peek_project_key(&mut buffer).await, project_key1);
 
         buffer
             .push(new_envelope(project_key3, None, None))
             .await
             .unwrap();
+        buffer.mark_ready(&project_key3, false);
+
         // All projects are not ready, so project 1 is on top (has the oldest envelopes):
-        assert_eq!(
-            buffer
-                .peek()
-                .await
-                .unwrap()
-                .envelope()
-                .unwrap()
-                .meta()
-                .public_key(),
-            project_key1
-        );
+        assert_eq!(peek_project_key(&mut buffer).await, project_key1);
 
         // After marking a project ready, it goes to the top:
         buffer.mark_ready(&project_key3, true);
-        assert_eq!(
-            buffer
-                .peek()
-                .await
-                .unwrap()
-                .envelope()
-                .unwrap()
-                .meta()
-                .public_key(),
-            project_key3
-        );
+        assert_eq!(peek_project_key(&mut buffer).await, project_key3);
         assert_eq!(
             buffer.pop().await.unwrap().unwrap().meta().public_key(),
             project_key3
         );
 
         // After popping, project 1 is on top again:
-        assert_eq!(
-            buffer
-                .peek()
-                .await
-                .unwrap()
-                .envelope()
-                .unwrap()
-                .meta()
-                .public_key(),
-            project_key1
-        );
+        assert_eq!(peek_project_key(&mut buffer).await, project_key1);
 
         // Mark project 1 as ready (still on top):
         buffer.mark_ready(&project_key1, true);
-        assert_eq!(
-            buffer
-                .peek()
-                .await
-                .unwrap()
-                .envelope()
-                .unwrap()
-                .meta()
-                .public_key(),
-            project_key1
-        );
+        assert_eq!(peek_project_key(&mut buffer).await, project_key1);
 
         // Mark project 2 as ready as well (now on top because most recent):
         buffer.mark_ready(&project_key2, true);
-        assert_eq!(
-            buffer
-                .peek()
-                .await
-                .unwrap()
-                .envelope()
-                .unwrap()
-                .meta()
-                .public_key(),
-            project_key2
-        );
+        assert_eq!(peek_project_key(&mut buffer).await, project_key2);
         assert_eq!(
             buffer.pop().await.unwrap().unwrap().meta().public_key(),
             project_key2
@@ -876,6 +825,9 @@ mod tests {
         let instant3 = envelope3.meta().start_time();
         buffer.push(envelope3).await.unwrap();
 
+        buffer.mark_ready(&project_key1, false);
+        buffer.mark_ready(&project_key2, false);
+
         // Nothing is ready, instant1 is on top:
         assert_eq!(
             buffer
@@ -1017,6 +969,9 @@ mod tests {
         buffer.push(envelope1).await.unwrap();
         buffer.push(envelope2).await.unwrap();
 
+        buffer.mark_ready(&project_key_1, false);
+        buffer.mark_ready(&project_key_2, false);
+
         // event_id_1 is first element:
         let Peek::NotReady(_, _, envelope) = buffer.peek().await.unwrap() else {
             panic!();