From c8353368a7e93d80808a165f2c2ebf0017204e49 Mon Sep 17 00:00:00 2001
From: Shankar Seal <sseal@microsoft.com>
Date: Sun, 17 Mar 2024 23:14:54 -0700
Subject: [PATCH] wait for all events.

---
 tests/api_test/api_test.cpp | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/tests/api_test/api_test.cpp b/tests/api_test/api_test.cpp
index 57b7f9d409..aa13b15cfa 100644
--- a/tests/api_test/api_test.cpp
+++ b/tests/api_test/api_test.cpp
@@ -27,6 +27,7 @@
 #include <ntsecapi.h>
 #include <thread>
 #include <vector>
+using namespace std::chrono_literals;
 
 CATCH_REGISTER_LISTENER(_watchdog)
 
@@ -1229,12 +1230,19 @@ TEST_CASE("ioctl_stress", "[stress]")
     bpf_object__close(object);
 }
 
+struct _ring_buffer_test_context
+{
+    uint32_t event_count;
+    uint32_t expected_event_count;
+    std::promise<void> promise;
+};
+
 TEST_CASE("test_ringbuffer_wraparound", "[stress]")
 {
     // Load bindmonitor_ringbuf.sys.
     struct bpf_object* object = nullptr;
     fd_t program_fd;
-    uint32_t event_count = 0;
+    struct _ring_buffer_test_context context;
     std::string app_id = "api_test.exe";
     uint32_t thread_count = 2;
     native_module_helper_t native_helper;
@@ -1254,14 +1262,21 @@ TEST_CASE("test_ringbuffer_wraparound", "[stress]")
     REQUIRE(max_iterations % thread_count == 0);
     uint32_t iterations_per_thread = max_iterations / thread_count;
 
+    // Initialize context.
+    context.event_count = 0;
+    context.expected_event_count = max_iterations;
+    auto ring_buffer_event_callback = context.promise.get_future();
     // Subscribe to the ring buffer.
     auto ring = ring_buffer__new(
         process_map_fd,
         [](void* ctx, void*, size_t) {
-            (*((uint32_t*)ctx))++;
+            struct _ring_buffer_test_context* context = reinterpret_cast<_ring_buffer_test_context*>(ctx);
+            if (++context->event_count == context->expected_event_count) {
+                context->promise.set_value();
+            }
             return 0;
         },
-        &event_count,
+        &context,
         nullptr);
 
     // Create 2 threads that invoke the program to trigger ring buffer events.
@@ -1290,7 +1305,8 @@ TEST_CASE("test_ringbuffer_wraparound", "[stress]")
     for (auto& t : threads) {
         t.join();
     }
-    REQUIRE(event_count == max_iterations);
+    // Wait for 1 second for the ring buffer to receive all events.
+    REQUIRE(ring_buffer_event_callback.wait_for(1s) == std::future_status::ready);
 
     // Unsubscribe from the ring buffer.
     ring_buffer__free(ring);