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 Dianzhang.Chen |
17 | * \date Dec 2020 |
18 | * \brief |
19 | */ |
20 | |
21 | #include "fake_collection.h" |
22 | #include "mock_collection_creator.h" |
23 | #include "port_helper.h" |
24 | #define private public |
25 | #define protected public |
26 | #include "repository/collection_manager.h" |
27 | #include "repository/repository_common/config.h" |
28 | #undef private |
29 | #undef protected |
30 | |
31 | #include "mock_index_agent_server.h" |
32 | |
33 | using namespace proxima::be::repository; |
34 | |
35 | static int PORT = 8010; |
36 | static int PID = 0; |
37 | |
38 | int main(int argc, char *argv[]) { |
39 | testing::InitGoogleTest(&argc, argv); |
40 | return RUN_ALL_TESTS(); |
41 | } |
42 | |
43 | class CollectionManagerTest : public ::testing::Test { |
44 | protected: |
45 | CollectionManagerTest() {} |
46 | |
47 | ~CollectionManagerTest() {} |
48 | void SetUp() { |
49 | PortHelper::GetPort(&PORT, &PID); |
50 | std::cout << "Server port: " << PORT << std::endl; |
51 | std::string index_uri = "127.0.0.1:" + std::to_string(PORT); |
52 | proxima::be::repository::Config::Instance() |
53 | .repository_config_.mutable_repository_config() |
54 | ->set_index_agent_addr(index_uri); |
55 | std::cout << "Set index addr: " << index_uri << std::endl; |
56 | } |
57 | void TearDown() { |
58 | PortHelper::RemovePortFile(PID); |
59 | } |
60 | }; |
61 | |
62 | TEST_F(CollectionManagerTest, TestGeneral) { |
63 | brpc::Server server_; |
64 | MockProximaServiceImpl svc_; |
65 | brpc::ServerOptions options; |
66 | ASSERT_EQ(0, server_.AddService(&svc_, brpc::SERVER_DOESNT_OWN_SERVICE)); |
67 | ASSERT_EQ(0, server_.Start(PORT, &options)); |
68 | { |
69 | MockCollectionCreatorPtr collection_creator = |
70 | std::make_shared<MockCollectionCreator>(); |
71 | EXPECT_CALL(*collection_creator, create(_)) |
72 | .WillRepeatedly(Invoke([](const CollectionInfo &info) -> CollectionPtr { |
73 | std::string collection_name = info.config().collection_name(); |
74 | MysqlHandlerPtr mysql_handler = |
75 | std::make_shared<MysqlHandler>(info.config()); |
76 | return std::make_shared<FakeMysqlCollection>(info.config(), |
77 | mysql_handler); |
78 | })); |
79 | |
80 | CollectionManagerPtr collection_manager = |
81 | std::make_shared<CollectionManager>(collection_creator); |
82 | int ret = collection_manager->init(); |
83 | ASSERT_EQ(0, ret); |
84 | |
85 | std::thread start_thread([&] { collection_manager->start(); }); |
86 | size_t count = 0; |
87 | auto created_collection = svc_.get_created_collections(); |
88 | auto collections_name = svc_.get_collections_name(); |
89 | while (count < 5) { // in order to sleep enough time |
90 | if (created_collection.size() == collections_name.size()) { |
91 | break; |
92 | } |
93 | std::this_thread::sleep_for(std::chrono::milliseconds(1000)); |
94 | created_collection = svc_.get_created_collections(); |
95 | collections_name = svc_.get_collections_name(); |
96 | count++; |
97 | } |
98 | EXPECT_EQ(created_collection.size(), 3); |
99 | ret = collection_manager->stop(); |
100 | sleep(2); |
101 | start_thread.join(); |
102 | EXPECT_EQ(0, ret); |
103 | EXPECT_EQ(created_collection.count("collection1" ), 1); |
104 | EXPECT_EQ(created_collection.count("collection2" ), 1); |
105 | EXPECT_EQ(created_collection.count("collection3" ), 1); |
106 | } |
107 | ASSERT_EQ(0, server_.Stop(0)); |
108 | ASSERT_EQ(0, server_.Join()); |
109 | } |
110 | |
111 | // todo<cdz>: Add update test when support update |
112 | // TEST_F(CollectionManagerTest, TestUpdate) { |
113 | // MockCollectionCreatorPtr collection_creator = |
114 | // std::make_shared<MockCollectionCreator>(); |
115 | // EXPECT_CALL(*collection_creator, create(_)) |
116 | // .WillRepeatedly( |
117 | // Invoke([](const CollectionConfig &config) -> CollectionPtr { |
118 | // std::string collection_name = config.collection_name(); |
119 | // MysqlHandlerPtr mysql_handler = |
120 | // std::make_shared<MysqlHandler>(config); |
121 | // return std::make_shared<FakeMysqlCollection>(config, |
122 | // mysql_handler); |
123 | // })); |
124 | |
125 | // CollectionManagerPtr collection_manager = |
126 | // std::make_shared<CollectionManager>(collection_creator); |
127 | // int ret = collection_manager->init(); |
128 | // ASSERT_EQ(0, ret); |
129 | |
130 | // CollectionInfo info1; |
131 | // CollectionInfo info2; |
132 | // CollectionInfo info3; |
133 | // info1.mutable_config()->set_collection_name("collection1"); |
134 | // info1.set_uuid("collection1-uuid"); |
135 | // info2.mutable_config()->set_collection_name("collection2"); |
136 | // info2.set_uuid("collection2-uuid"); |
137 | // info3.mutable_config()->set_collection_name("collection3"); |
138 | // info3.set_uuid("collection3-uuid"); |
139 | // std::vector<CollectionInfo> collection_infos{info1, info2, info3}; |
140 | // std::vector<CollectionInfo> new_collections; |
141 | // std::vector<std::string> old_collections; |
142 | // std::vector<std::string> expired_collections; |
143 | // collection_manager->classify_collections(collection_infos, |
144 | // &new_collections, |
145 | // &old_collections, |
146 | // &expired_collections); |
147 | // EXPECT_EQ(3, new_collections.size()); |
148 | // EXPECT_EQ(0, old_collections.size()); |
149 | // EXPECT_EQ(0, expired_collections.size()); |
150 | // EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(), |
151 | // [](const CollectionInfo &info) { |
152 | // return info.config().collection_name() == |
153 | // "collection1" && |
154 | // info.uuid() == |
155 | // "collection1-uuid"; |
156 | // }), |
157 | // true); |
158 | // EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(), |
159 | // [](const CollectionInfo &info) { |
160 | // return info.config().collection_name() == |
161 | // "collection2" && |
162 | // info.uuid() == |
163 | // "collection2-uuid"; |
164 | // }), |
165 | // true); |
166 | // EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(), |
167 | // [](const CollectionInfo &info) { |
168 | // return info.config().collection_name() == |
169 | // "collection3" && |
170 | // info.uuid() == |
171 | // "collection3-uuid"; |
172 | // }), |
173 | // true); |
174 | |
175 | // // collection_manager->create_collections(new_collections); |
176 | // std::thread worker( |
177 | // [&]() { collection_manager->create_collections(new_collections); }); |
178 | // size_t count = 0; |
179 | // auto created_collection = svc_.get_created_collections(); |
180 | // auto collections_name = svc_.get_collections_name(); |
181 | // while (count < 5) { // in order to sleep enough time |
182 | // if (created_collection.size() == collections_name.size()) { |
183 | // break; |
184 | // } |
185 | // std::this_thread::sleep_for(std::chrono::milliseconds(1000)); |
186 | // created_collection = svc_.get_created_collections(); |
187 | // collections_name = svc_.get_collections_name(); |
188 | // count++; |
189 | // } |
190 | // EXPECT_EQ(created_collection.size(), |
191 | // collection_manager->collections_.size()); |
192 | // collection_manager->update_collections(old_collections); |
193 | // collection_manager->drop_collections(expired_collections); |
194 | // EXPECT_EQ(3, collection_manager->collections_.size()); |
195 | // EXPECT_EQ(1, collection_manager->collections_.count("collection1-uuid")); |
196 | // EXPECT_EQ(1, collection_manager->collections_.count("collection2-uuid")); |
197 | // EXPECT_EQ(1, collection_manager->collections_.count("collection3-uuid")); |
198 | |
199 | // // Update |
200 | // collections_config[0].set_schema_revision(1); |
201 | // collections_config[1].set_schema_revision(2); |
202 | // collections_config[2].set_schema_revision(2); |
203 | // new_collections.clear(); |
204 | // old_collections.clear(); |
205 | // expired_collections.clear(); |
206 | // collection_manager->classify_collections(collections_config, |
207 | // &new_collections, |
208 | // &old_collections, |
209 | // &expired_collections); |
210 | // EXPECT_EQ(0, new_collections.size()); |
211 | // EXPECT_EQ(2, old_collections.size()); |
212 | // EXPECT_EQ(0, expired_collections.size()); |
213 | // EXPECT_EQ(std::any_of(old_collections.begin(), old_collections.end(), |
214 | // [](const std::string &name) { |
215 | // return name == "collection2-uuid"; |
216 | // }), |
217 | // true); |
218 | // EXPECT_EQ(std::any_of(old_collections.begin(), old_collections.end(), |
219 | // [](const std::string &name) { |
220 | // return name == "collection3-uuid"; |
221 | // }), |
222 | // true); |
223 | // collection_manager->create_collections(new_collections); |
224 | // collection_manager->update_collections(old_collections); |
225 | // collection_manager->drop_collections(expired_collections); |
226 | // EXPECT_EQ(1, collection_manager->collections_.count("collection1-uuid")); |
227 | // EXPECT_EQ(1, collection_manager->collections_.count("collection2-uuid")); |
228 | // EXPECT_EQ(1, collection_manager->collections_.count("collection3-uuid")); |
229 | // EXPECT_EQ( |
230 | // 1, |
231 | // collection_manager->collections_["collection1-uuid"]->schema_revision()); |
232 | // EXPECT_EQ( |
233 | // 2, |
234 | // collection_manager->collections_["collection2-uuid"]->schema_revision()); |
235 | // EXPECT_EQ( |
236 | // 2, |
237 | // collection_manager->collections_["collection3-uuid"]->schema_revision()); |
238 | // ret = collection_manager->stop(); |
239 | // EXPECT_EQ(ret, 0); |
240 | // worker.join(); |
241 | |
242 | // ASSERT_EQ(0, server_.Stop(0)); |
243 | // ASSERT_EQ(0, server_.Join()); |
244 | // } |
245 | |