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
31using GetDocumentRequest = proxima::be::proto::GetDocumentRequest;
32using GetDocumentResponse = proxima::be::proto::GetDocumentResponse;
33
34class 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
80TEST_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
90TEST_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
148TEST_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
166TEST_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
290TEST_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