1 | /** |
2 | * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | * |
16 | * \author guonix |
17 | * \date Dec 2020 |
18 | * \brief |
19 | */ |
20 | |
21 | |
22 | #include "query/equal_query.h" |
23 | #include <gtest/gtest.h> |
24 | #include <meta/meta_impl.h> |
25 | #include "index/mock_index_service.h" // for MockIndexService |
26 | #include "index/mock_segment.h" // for MockSegment |
27 | #include "meta/mock_meta_service.h" // for MockMetaService |
28 | #include "mock_executor.h" // for MockExecutor |
29 | |
30 | |
31 | using GetDocumentRequest = proxima::be::proto::GetDocumentRequest; |
32 | using GetDocumentResponse = proxima::be::proto::GetDocumentResponse; |
33 | |
34 | class EqualQueryTest : public Test { |
35 | protected: |
36 | // Sets up the test fixture. |
37 | void SetUp() override { |
38 | init_request(); |
39 | init_response(); |
40 | } |
41 | |
42 | // Tears down the test fixture. |
43 | void TearDown() override { |
44 | cleanup_request(); |
45 | cleanup_response(); |
46 | } |
47 | |
48 | private: |
49 | void init_request() { |
50 | request_ = new (std::nothrow) GetDocumentRequest(); |
51 | request_->set_collection_name(collection_); |
52 | request_->set_debug_mode(false); |
53 | request_->set_primary_key(1); |
54 | } |
55 | |
56 | void init_response() { |
57 | response_ = new (std::nothrow) GetDocumentResponse(); |
58 | } |
59 | |
60 | template <typename Pointer> |
61 | void delete_pointer_if(Pointer *ptr) { |
62 | delete ptr; |
63 | ptr = nullptr; |
64 | } |
65 | |
66 | void cleanup_request() { |
67 | delete_pointer_if(request_); |
68 | } |
69 | |
70 | void cleanup_response() { |
71 | delete_pointer_if(response_); |
72 | } |
73 | |
74 | protected: |
75 | GetDocumentRequest *request_{nullptr}; |
76 | GetDocumentResponse *response_{nullptr}; |
77 | std::string collection_{"unittest" }; |
78 | }; |
79 | |
80 | TEST_F(EqualQueryTest, TestBaseFunctional) { |
81 | // Test invalid params |
82 | auto query = std::make_shared<EqualQuery>(0, nullptr, nullptr, nullptr, |
83 | nullptr, nullptr, nullptr); |
84 | |
85 | EXPECT_EQ(query->mode(), IOMode::READONLY); |
86 | EXPECT_EQ(query->type(), QueryType::EQUAL); |
87 | EXPECT_EQ(query->id(), 0); |
88 | } |
89 | |
90 | TEST_F(EqualQueryTest, TestValidate) { |
91 | auto executor = std::make_shared<MockExecutor>(); |
92 | auto meta_service = std::make_shared<MockMetaService>(); |
93 | auto index_service = std::make_shared<MockIndexService>(); |
94 | auto meta = std::make_shared<MetaWrapper>(meta_service); |
95 | // Test invalid params |
96 | auto query = std::make_shared<EqualQuery>( |
97 | 0, nullptr, index_service, meta, executor, |
98 | std::make_shared<proxima::be::Profiler>(false), response_); |
99 | EXPECT_TRUE(query->validate() != 0); |
100 | |
101 | query.reset(new (std::nothrow) EqualQuery( |
102 | 0, request_, index_service, meta, nullptr, |
103 | std::make_shared<proxima::be::Profiler>(false), response_)); |
104 | EXPECT_TRUE(query->validate() != 0); |
105 | |
106 | // Return 0, with invalid collection meta |
107 | EXPECT_CALL(*meta_service, get_current_collection(_)) |
108 | .WillOnce(Invoke( |
109 | [](const std::string &) -> CollectionMetaPtr { return nullptr; })) |
110 | .RetiresOnSaturation(); // success |
111 | |
112 | query.reset(new (std::nothrow) EqualQuery( |
113 | 0, request_, index_service, meta, executor, |
114 | std::make_shared<proxima::be::Profiler>(false), nullptr)); |
115 | EXPECT_TRUE(query->validate() != 0); |
116 | |
117 | CollectionMeta collection_meta; |
118 | collection_meta.mutable_forward_columns()->push_back("forward1" ); |
119 | collection_meta.mutable_forward_columns()->push_back("forward2" ); |
120 | auto column1 = std::make_shared<ColumnMeta>(); |
121 | collection_meta.append(column1); |
122 | |
123 | CollectionImplPtr collection = |
124 | std::make_shared<CollectionImpl>(collection_meta); |
125 | |
126 | // Return 0, with invalid collection meta |
127 | EXPECT_CALL(*meta_service, get_current_collection(_)) |
128 | .WillOnce(Invoke( |
129 | [](const std::string &) -> CollectionMetaPtr { return nullptr; })) |
130 | .RetiresOnSaturation(); // success |
131 | // Set all right arguments |
132 | query.reset(new (std::nothrow) EqualQuery( |
133 | 0, request_, index_service, meta, executor, |
134 | std::make_shared<proxima::be::Profiler>(false), response_)); |
135 | // Can't valid column from meta wrapper |
136 | EXPECT_TRUE(query->validate() != 0); |
137 | |
138 | // Return collection |
139 | EXPECT_CALL(*meta_service, get_current_collection("unittest" )) |
140 | .WillOnce(Invoke([&collection](const std::string &) -> CollectionMetaPtr { |
141 | return collection->meta(); |
142 | })) |
143 | .RetiresOnSaturation(); // success |
144 | |
145 | EXPECT_EQ(query->validate(), 0); |
146 | } |
147 | |
148 | TEST_F(EqualQueryTest, TestPrepare) { |
149 | auto executor = std::make_shared<MockExecutor>(); |
150 | auto meta_service = std::make_shared<MockMetaService>(); |
151 | auto index_service = std::make_shared<MockIndexService>(); |
152 | auto meta = std::make_shared<MetaWrapper>(meta_service); |
153 | |
154 | EXPECT_CALL(*index_service, list_segments(collection_, _)) |
155 | .WillOnce(Return(1)) |
156 | .WillOnce(Return(0)) // Success but no available segments |
157 | .RetiresOnSaturation(); |
158 | |
159 | auto query = std::make_shared<EqualQuery>( |
160 | 0, request_, index_service, meta, executor, |
161 | std::make_shared<proxima::be::Profiler>(false), response_); |
162 | EXPECT_TRUE(query->prepare() != 0); |
163 | EXPECT_TRUE(query->prepare() != 0); |
164 | } |
165 | |
166 | TEST_F(EqualQueryTest, TestEvaluate) { |
167 | auto executor = std::make_shared<MockExecutor>(); |
168 | auto meta_service = std::make_shared<MockMetaService>(); |
169 | auto index_service = std::make_shared<MockIndexService>(); |
170 | auto meta = std::make_shared<MetaWrapper>(meta_service); |
171 | auto segment = std::make_shared<MockSegment>(); |
172 | |
173 | EXPECT_CALL(*index_service, list_segments(_, _)) |
174 | .WillOnce(Invoke([&segment](const std::string &, |
175 | index::SegmentPtrList *segments) -> int { |
176 | segments->push_back(segment); |
177 | return 0; |
178 | })) |
179 | .RetiresOnSaturation(); |
180 | |
181 | CollectionImplPtr collection_impl = nullptr; |
182 | { // Init collection |
183 | CollectionMeta cmeta; |
184 | cmeta.mutable_forward_columns()->push_back("forward1" ); |
185 | cmeta.mutable_forward_columns()->push_back("forward2" ); |
186 | cmeta.set_revision(10); |
187 | collection_impl.reset(new CollectionImpl(cmeta)); |
188 | } |
189 | |
190 | EXPECT_CALL(*meta_service, get_collection(_, _)) |
191 | .WillRepeatedly( |
192 | Invoke([this, &collection_impl](const std::string &collection, |
193 | uint64_t revision) { |
194 | EXPECT_EQ(collection_, collection); |
195 | EXPECT_EQ(revision, 10); |
196 | return collection_impl->meta(); |
197 | })) |
198 | .RetiresOnSaturation(); |
199 | |
200 | auto query = std::make_shared<EqualQuery>( |
201 | 0, request_, index_service, meta, executor, |
202 | std::make_shared<proxima::be::Profiler>(false), response_); |
203 | |
204 | EXPECT_EQ(query->prepare(), 0); |
205 | |
206 | { // evaluate failed with fake execute |
207 | EXPECT_CALL(*executor, execute_tasks(_)) |
208 | .WillOnce(Return(0)) |
209 | .RetiresOnSaturation(); |
210 | |
211 | // Success, but with no results |
212 | EXPECT_EQ(query->evaluate(), 0); |
213 | } |
214 | |
215 | { // Evaluate success, but no enough values |
216 | // Execute task |
217 | EXPECT_CALL(*executor, execute_tasks(_)) |
218 | .WillOnce(Invoke([](const TaskPtrList &tasks) { |
219 | for (auto &task : tasks) { |
220 | task->status(Task::Status::SCHEDULED); |
221 | task->run(); |
222 | } |
223 | return 0; |
224 | })) // Fake Execute |
225 | .RetiresOnSaturation(); |
226 | |
227 | // Set results |
228 | EXPECT_CALL(*segment, kv_search(_, _)) |
229 | .WillOnce(Invoke([](uint64_t primary_key, QueryResult *result) { |
230 | EXPECT_EQ(primary_key, 1); |
231 | result->primary_key = 1; |
232 | result->revision = 10; |
233 | return 0; |
234 | })) |
235 | .RetiresOnSaturation(); |
236 | |
237 | // Expect mismatch error |
238 | EXPECT_EQ(query->evaluate(), PROXIMA_BE_ERROR_CODE(MismatchedForward)); |
239 | } |
240 | |
241 | response_->Clear(); |
242 | |
243 | { // Test serialize |
244 | EXPECT_CALL(*executor, execute_tasks(_)) |
245 | .WillOnce(Invoke([](const TaskPtrList &tasks) { |
246 | for (auto &task : tasks) { |
247 | task->status(Task::Status::SCHEDULED); |
248 | task->run(); |
249 | } |
250 | return 0; |
251 | })) // Fake Execute |
252 | .RetiresOnSaturation(); |
253 | |
254 | // Set results |
255 | EXPECT_CALL(*segment, kv_search(_, _)) |
256 | .WillOnce(Invoke([](uint64_t primary_key, QueryResult *result) { |
257 | EXPECT_EQ(primary_key, 1); |
258 | result->primary_key = 1U; |
259 | result->lsn = 1U; |
260 | result->revision = 10; |
261 | result->score = 0.95f; |
262 | proxima::be::proto::GenericValueList values; |
263 | auto value = values.add_values(); |
264 | value->set_int32_value(10); |
265 | value = values.add_values(); |
266 | value->set_string_value("strvalue" ); |
267 | // Forward |
268 | result->forward_data.assign(values.SerializeAsString()); |
269 | return 0; |
270 | })) |
271 | .RetiresOnSaturation(); |
272 | |
273 | // Expect 0 |
274 | EXPECT_EQ(query->evaluate(), 0); |
275 | |
276 | |
277 | EXPECT_EQ(response_->document().primary_key(), 1U); |
278 | EXPECT_EQ(response_->document().forward_column_values_size(), 2); |
279 | auto &kv = response_->document().forward_column_values(0); |
280 | EXPECT_EQ(kv.key(), "forward1" ); |
281 | EXPECT_EQ(kv.value().int32_value(), 10); |
282 | auto &kv1 = response_->document().forward_column_values(1); |
283 | EXPECT_EQ(kv1.key(), "forward2" ); |
284 | EXPECT_EQ(kv1.value().string_value(), "strvalue" ); |
285 | } |
286 | |
287 | response_->Clear(); |
288 | } |
289 | |
290 | TEST_F(EqualQueryTest, TestFinalize) { |
291 | auto executor = std::make_shared<MockExecutor>(); |
292 | auto meta_service = std::make_shared<MockMetaService>(); |
293 | auto index_service = std::make_shared<MockIndexService>(); |
294 | auto meta = std::make_shared<MetaWrapper>(meta_service); |
295 | auto query = std::make_shared<EqualQuery>( |
296 | 0, request_, index_service, meta, executor, |
297 | std::make_shared<proxima::be::Profiler>(false), response_); |
298 | |
299 | EXPECT_EQ(query->finalize(), 0); |
300 | } |
301 | |