1#include <gtest/gtest.h>
2
3#include <Core/Types.h>
4#include <IO/ReadHelpers.h>
5#include <IO/ReadBufferFromString.h>
6#include <IO/ConcatReadBuffer.h>
7#include <IO/PeekableReadBuffer.h>
8
9namespace DB::ErrorCodes
10{
11 extern const int LOGICAL_ERROR;
12 extern const int MEMORY_LIMIT_EXCEEDED;
13}
14
15static void readAndAssert(DB::ReadBuffer & buf, const char * str)
16{
17 size_t n = strlen(str);
18 char tmp[n];
19 buf.readStrict(tmp, n);
20 ASSERT_EQ(strncmp(tmp, str, n), 0);
21}
22
23static void assertAvailable(DB::ReadBuffer & buf, const char * str)
24{
25 size_t n = strlen(str);
26 ASSERT_EQ(buf.available(), n);
27 ASSERT_EQ(strncmp(buf.position(), str, n), 0);
28}
29
30TEST(PeekableReadBuffer, CheckpointsWorkCorrectly)
31try
32{
33 std::string s1 = "0123456789";
34 std::string s2 = "qwertyuiop";
35 std::string s3 = "asdfghjkl;";
36 std::string s4 = "zxcvbnm,./";
37 DB::ReadBufferFromString b1(s1);
38 DB::ReadBufferFromString b2(s2);
39 DB::ReadBufferFromString b3(s3);
40 DB::ReadBufferFromString b4(s4);
41
42 DB::ConcatReadBuffer concat({&b1, &b2, &b3, &b4});
43 DB::PeekableReadBuffer peekable(concat, 0, 16);
44
45 ASSERT_TRUE(!peekable.eof());
46 assertAvailable(peekable, "0123456789");
47 {
48 DB::PeekableReadBufferCheckpoint checkpoint{peekable};
49 readAndAssert(peekable, "01234");
50 }
51 bool exception = false;
52 try
53 {
54 peekable.rollbackToCheckpoint();
55 }
56 catch (DB::Exception & e)
57 {
58 if (e.code() != DB::ErrorCodes::LOGICAL_ERROR)
59 throw;
60 exception = true;
61 }
62 ASSERT_TRUE(exception);
63 assertAvailable(peekable, "56789");
64
65 readAndAssert(peekable, "56");
66
67 peekable.setCheckpoint();
68 readAndAssert(peekable, "789qwertyu");
69 peekable.rollbackToCheckpoint();
70 peekable.dropCheckpoint();
71 assertAvailable(peekable, "789");
72
73 exception = false;
74 try
75 {
76 DB::PeekableReadBufferCheckpoint checkpoint{peekable, true};
77 peekable.ignore(30);
78 }
79 catch (DB::Exception & e)
80 {
81 if (e.code() != DB::ErrorCodes::MEMORY_LIMIT_EXCEEDED)
82 throw;
83 exception = true;
84 }
85 ASSERT_TRUE(exception);
86 assertAvailable(peekable, "789qwertyuiop");
87
88 readAndAssert(peekable, "789qwertyu");
89 peekable.setCheckpoint();
90 readAndAssert(peekable, "iopasdfghj");
91 assertAvailable(peekable, "kl;");
92 peekable.dropCheckpoint();
93
94 peekable.setCheckpoint();
95 readAndAssert(peekable, "kl;zxcvbnm,./");
96 ASSERT_TRUE(peekable.eof());
97 ASSERT_TRUE(peekable.eof());
98 ASSERT_TRUE(peekable.eof());
99 peekable.rollbackToCheckpoint();
100 readAndAssert(peekable, "kl;zxcvbnm");
101 peekable.dropCheckpoint();
102
103 ASSERT_TRUE(peekable.hasUnreadData());
104 readAndAssert(peekable, ",./");
105 ASSERT_FALSE(peekable.hasUnreadData());
106
107 ASSERT_TRUE(peekable.eof());
108 ASSERT_TRUE(peekable.eof());
109 ASSERT_TRUE(peekable.eof());
110
111}
112catch (const DB::Exception & e)
113{
114 std::cerr << e.what() << ", " << e.displayText() << std::endl;
115 throw;
116}
117
118