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 | |
9 | namespace DB::ErrorCodes |
10 | { |
11 | extern const int LOGICAL_ERROR; |
12 | extern const int MEMORY_LIMIT_EXCEEDED; |
13 | } |
14 | |
15 | static 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 | |
23 | static 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 | |
30 | TEST(PeekableReadBuffer, CheckpointsWorkCorrectly) |
31 | try |
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 | } |
112 | catch (const DB::Exception & e) |
113 | { |
114 | std::cerr << e.what() << ", " << e.displayText() << std::endl; |
115 | throw; |
116 | } |
117 | |
118 | |