1 files changed, 111 insertions, 0 deletions
diff --git a/src/seastar/tests/unit/stream_reader_test.cc b/src/seastar/tests/unit/stream_reader_test.cc
new file mode 100644
index 000000000..61401e08e
--- /dev/null
+++ b/src/seastar/tests/unit/stream_reader_test.cc
@@ -0,0 +1,111 @@
+/*
+ * This file is open source software, licensed to you under the terms
+ * of the Apache License, Version 2.0 (the "License").  See the NOTICE file
+ * distributed with this work for additional information regarding copyright
+ * ownership.  You may not use this file except in compliance with the License.
+ *
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*
+ * Copyright (C) 2020 ScyllaDB.
+ */
+
+#include <seastar/core/future.hh>
+#include <seastar/core/temporary_buffer.hh>
+#include <seastar/core/thread.hh>
+#include <seastar/testing/test_case.hh>
+#include <seastar/testing/thread_test_case.hh>
+#include <seastar/http/request.hh>
+#include <seastar/util/short_streams.hh>
+#include <string>
+
+using namespace seastar;
+using namespace util;
+
+/*
+ * Simple data source producing up to total_size bytes
+ * in buffer_size-byte chunks.
+ * */
+class test_source_impl : public data_source_impl {
+    short _current_letter = 0; // a-z corresponds to 0-25
+    size_t _buffer_size;
+    size_t _remaining_size;
+public:
+    test_source_impl(size_t buffer_size, size_t total_size)
+        : _buffer_size(buffer_size), _remaining_size(total_size) {
+    }
+    virtual future<temporary_buffer<char>> get() override {
+        size_t len = std::min(_buffer_size, _remaining_size);
+        temporary_buffer<char> tmp(len);
+        for (size_t i = 0; i < len; i++) {
+            tmp.get_write()[i] = 'a' + _current_letter;
+            ++_current_letter %= 26;
+        }
+        _remaining_size -= len;
+        return make_ready_future<temporary_buffer<char>>(std::move(tmp));
+    }
+    virtual future<temporary_buffer<char>> skip(uint64_t n) override {
+        _remaining_size -= std::min(_remaining_size, n);
+        _current_letter += n %= 26;
+        return make_ready_future<temporary_buffer<char>>();
+    }
+};
+
+SEASTAR_TEST_CASE(test_read_all) {
+    return async([] {
+        auto check_read_all = [] (input_stream<char>& strm, const char* test) {
+            auto all = read_entire_stream(strm).get0();
+            sstring s;
+            for (auto&& buf: all) {
+                s += seastar::to_sstring(std::move(buf));
+            };
+            BOOST_REQUIRE_EQUAL(s, test);
+        };
+        input_stream<char> inp(data_source(std::make_unique<test_source_impl>(5, 15)));
+        check_read_all(inp, "abcdefghijklmno");
+        BOOST_REQUIRE(inp.eof());
+        input_stream<char> inp2(data_source(std::make_unique<test_source_impl>(5, 16)));
+        check_read_all(inp2, "abcdefghijklmnop");
+        BOOST_REQUIRE(inp2.eof());
+        input_stream<char> empty_inp(data_source(std::make_unique<test_source_impl>(5, 0)));
+        check_read_all(empty_inp, "");
+        BOOST_REQUIRE(empty_inp.eof());
+
+        input_stream<char> inp_cont(data_source(std::make_unique<test_source_impl>(5, 15)));
+        BOOST_REQUIRE_EQUAL(to_sstring(read_entire_stream_contiguous(inp_cont).get0()), "abcdefghijklmno");
+        BOOST_REQUIRE(inp_cont.eof());
+        input_stream<char> inp_cont2(data_source(std::make_unique<test_source_impl>(5, 16)));
+        BOOST_REQUIRE_EQUAL(to_sstring(read_entire_stream_contiguous(inp_cont2).get0()), "abcdefghijklmnop");
+        BOOST_REQUIRE(inp_cont2.eof());
+        input_stream<char> empty_inp_cont(data_source(std::make_unique<test_source_impl>(5, 0)));
+        BOOST_REQUIRE_EQUAL(to_sstring(read_entire_stream_contiguous(empty_inp_cont).get0()), "");
+        BOOST_REQUIRE(empty_inp_cont.eof());
+    });
+}
+
+SEASTAR_TEST_CASE(test_skip_all) {
+    return async([] {
+        input_stream<char> inp(data_source(std::make_unique<test_source_impl>(5, 15)));
+        skip_entire_stream(inp).get();
+        BOOST_REQUIRE(inp.eof());
+        BOOST_REQUIRE(to_sstring(inp.read().get0()).empty());
+        input_stream<char> inp2(data_source(std::make_unique<test_source_impl>(5, 16)));
+        skip_entire_stream(inp2).get();
+        BOOST_REQUIRE(inp2.eof());
+        BOOST_REQUIRE(to_sstring(inp2.read().get0()).empty());
+        input_stream<char> empty_inp(data_source(std::make_unique<test_source_impl>(5, 0)));
+        skip_entire_stream(empty_inp).get();
+        BOOST_REQUIRE(empty_inp.eof());
+        BOOST_REQUIRE(to_sstring(empty_inp.read().get0()).empty());
+    });
+}