From 470beda9f047c5cad55f0a1cc96ec7c1127cc07b Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Fri, 4 Feb 2022 16:55:29 -0500
Subject: [PATCH] Add nvbench::state stream tests

---
 testing/state.cu | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/testing/state.cu b/testing/state.cu
index 58e8dd1..ade63ed 100644
--- a/testing/state.cu
+++ b/testing/state.cu
@@ -51,6 +51,23 @@ struct state_tester : public nvbench::state
 
 using nvbench::detail::state_tester;
 
+void test_streams()
+{
+  dummy_bench bench;
+
+  state_tester state{bench};
+
+  // Test non-owning stream
+  state.set_cuda_stream(nvbench::cuda_stream{cudaStreamDefault, false});
+  ASSERT(state.get_cuda_stream() == cudaStreamDefault);
+
+  // Test owning stream
+  auto stream = nvbench::cuda_stream{};
+  auto gold   = stream.get_stream();
+  state.set_cuda_stream(std::move(stream));
+  ASSERT(state.get_cuda_stream() == gold);
+}
+
 void test_params()
 {
   dummy_bench bench;
@@ -110,6 +127,7 @@ void test_defaults()
 
 int main()
 {
+  test_streams();
   test_params();
   test_summaries();
   test_defaults();