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
33using namespace proxima::be::repository;
34
35static int PORT = 8010;
36static int PID = 0;
37
38int main(int argc, char *argv[]) {
39 testing::InitGoogleTest(&argc, argv);
40 return RUN_ALL_TESTS();
41}
42
43class 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
62TEST_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