Add some more async_stream tests.

lewissbaker · lewissbaker · commit a36dcdbf2ee7 · 2018-01-25T16:42:06.000+10:30
- Test for terminating consumption of a stream before the end.
- Test for composing streams using take(n, subscribable)
diff --git a/include/cppcoro/async_stream.hpp b/include/cppcoro/async_stream.hpp
@@ -320,6 +320,8 @@ namespace cppcoro
 
 	public:
 
+		using value_type = std::remove_reference_t<T>;
+
 		async_stream(handle_t coroutine) noexcept
 			: m_coroutine(coroutine)
 		{}
diff --git a/include/cppcoro/consume.hpp b/include/cppcoro/consume.hpp
@@ -5,6 +5,8 @@
 #ifndef CPPCORO_CONSUME_HPP_INCLUDED
 #define CPPCORO_CONSUME_HPP_INCLUDED
 
+#include <functional>
+
 namespace cppcoro
 {
 	template<typename SUBSCRIBABLE, typename CONSUME_FUNC>
diff --git a/test/async_stream_tests.cpp b/test/async_stream_tests.cpp
@@ -43,4 +43,100 @@ TEST_CASE("consume async_stream")
 	CHECK(result == 3);
 }
 
+TEST_CASE("producer exiting early on destruction of stream")
+{
+	int lastProduced = -1;
+	bool ranFinalisation = false;
+
+	auto subscribable = make_subscribable(
+		[&]() -> async_stream_subscription<int>
+		{
+			for (int i = 0; i < 5; ++i)
+			{
+				lastProduced = i;
+				bool produceMore = co_yield i;
+				if (!produceMore) break;
+			}
+
+			ranFinalisation = true;
+		});
+
+	sync_wait(consume(subscribable, [](async_stream<int> stream) -> task<>
+	{
+		// TRICKY: Need to move parameter to local var here to ensure it is destructed when the
+		// coroutine runs to completion. Normally, the destructor of parameters won't run until
+		// the coroutine is destroyed. However, the coroutine won't be destroyed until the
+		// entire consume operation completes, which includes waiting until the producer coroutine
+		// runs to completion. However, the producer coroutine isn't going to run to completion
+		// unless we either consume it completely or detach from the stream, causing a deadlock.
+		auto localStream = std::move(stream);
+		for co_await(int value : localStream)
+		{
+			// Don't consume any more values once we get to '3'.
+			if (value == 3) break;
+		}
+	}));
+
+	CHECK(lastProduced == 3);
+	CHECK(ranFinalisation);
+}
+
+template<typename COUNT, typename SUBSCRIBABLE>
+auto take(COUNT n, SUBSCRIBABLE&& s)
+{
+	return make_subscribable(
+		[n, s = std::forward<SUBSCRIBABLE>(s)]
+		{
+			auto[sourceStream, sourceTask] = s.subscribe();
+
+		using value_type = typename std::remove_reference_t<decltype(sourceStream)>::value_type;
+
+		auto[outputStream, outputTask] = [](COUNT n, auto stream) -> async_stream_subscription<value_type>
+		{
+			auto localStream = std::move(stream);
+			if (n > 0)
+			{
+				for co_await(auto&& value : localStream)
+				{
+					bool produceMore = co_yield value;
+					if (!produceMore) break;
+					if (--n == 0) break;
+				}
+			}
+		}(n, std::move(sourceStream));
+
+		return std::make_tuple(
+			std::move(outputStream),
+			when_all(std::move(sourceTask), std::move(outputTask)) | fmap([](auto) {}));
+		});
+}
+
+TEST_CASE("take(5)")
+{
+	auto subscribable = take(5, make_subscribable([]() -> async_stream_subscription<int>
+	{
+		for (int i = 0; i < 10; ++i)
+		{
+			bool produceMore = co_yield i;
+			if (!produceMore) break;
+		}
+	}));
+
+	auto values = sync_wait(consume(subscribable, [](async_stream<int> stream) -> task<std::vector<int>>
+	{
+		auto localStream = std::move(stream);
+
+		std::vector<int> values;
+		for co_await(int value : localStream)
+		{
+			values.push_back(value);
+		}
+		return std::move(values);
+	}));
+
+	CHECK(values.size() == 5);
+	CHECK(values[0] == 0);
+	CHECK(values[4] == 4);
+}
+
 TEST_SUITE_END();

Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,8 @@`
`5`	`5`	`#ifndef CPPCORO_CONSUME_HPP_INCLUDED`
`6`	`6`	`#define CPPCORO_CONSUME_HPP_INCLUDED`
`7`	`7`
	`8`	`+#include <functional>`
	`9`	`+`
`8`	`10`	`namespace cppcoro`
`9`	`11`	`{`
`10`	`12`	`template<typename SUBSCRIBABLE, typename CONSUME_FUNC>`