1 files changed, 23 insertions, 0 deletions
diff --git a/test/xray/TestCases/Posix/logging-modes.cc b/test/xray/TestCases/Posix/logging-modes.cc
index 22f6942b7..b846cf675 100644
--- a/test/xray/TestCases/Posix/logging-modes.cc
+++ b/test/xray/TestCases/Posix/logging-modes.cc
@@ -9,6 +9,7 @@
 #include "xray/xray_log_interface.h"
 #include <cassert>
 #include <cstdio>
+#include <string>
 
 [[clang::xray_never_instrument]] void printing_handler(int32_t fid,
                                                        XRayEntryType) {
@@ -20,8 +21,19 @@
   printing = false;
 }
 
+[[clang::xray_never_instrument]] XRayBuffer next_buffer(XRayBuffer buffer) {
+  static const char data[10] = {};
+  static const XRayBuffer first_and_last{data, 10};
+  if (buffer.Data == nullptr)
+    return first_and_last;
+  if (buffer.Data == first_and_last.Data)
+    return XRayBuffer{nullptr, 0};
+  assert(false && "Invalid buffer provided.");
+}
+
 [[clang::xray_never_instrument]] XRayLogInitStatus
 printing_init(size_t, size_t, void *, size_t) {
+  __xray_log_set_buffer_iterator(next_buffer);
   return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
 }
 
@@ -30,11 +42,16 @@ printing_init(size_t, size_t, void *, size_t) {
 }
 
 [[clang::xray_never_instrument]] XRayLogFlushStatus printing_flush_log() {
+  __xray_log_remove_buffer_iterator();
   return XRayLogFlushStatus::XRAY_LOG_FLUSHED;
 }
 
 [[clang::xray_always_instrument]] void callme() { std::printf("called me!\n"); }
 
+static auto buffer_counter = 0;
+
+void process_buffer(const char *, XRayBuffer) { ++buffer_counter; }
+
 static bool unused = [] {
   assert(__xray_log_register_mode("custom",
                                   {printing_init, printing_finalize,
@@ -46,6 +63,9 @@ static bool unused = [] {
 int main(int argc, char **argv) {
   assert(__xray_log_select_mode("custom") ==
          XRayLogRegisterStatus::XRAY_REGISTRATION_OK);
+  assert(__xray_log_get_current_mode() != nullptr);
+  std::string current_mode = __xray_log_get_current_mode();
+  assert(current_mode == "custom");
   assert(__xray_patch() == XRayPatchingStatus::SUCCESS);
   assert(__xray_log_init(0, 0, nullptr, 0) ==
          XRayLogInitStatus::XRAY_LOG_INITIALIZED);
@@ -53,6 +73,9 @@ int main(int argc, char **argv) {
   callme(); // CHECK: called me!
   // CHECK: printing {{.*}}
   assert(__xray_log_finalize() == XRayLogInitStatus::XRAY_LOG_FINALIZED);
+  assert(__xray_log_process_buffers(process_buffer) ==
+         XRayLogFlushStatus::XRAY_LOG_FLUSHED);
+  assert(buffer_counter == 1);
   assert(__xray_log_flushLog() == XRayLogFlushStatus::XRAY_LOG_FLUSHED);
   assert(__xray_log_select_mode("not-found") ==
          XRayLogRegisterStatus::XRAY_MODE_NOT_FOUND);