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 CollectionManagerDropTest : public ::testing::Test { |
44 | protected: |
45 | CollectionManagerDropTest() {} |
46 | |
47 | ~CollectionManagerDropTest() {} |
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(CollectionManagerDropTest, TestDrop) { |
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 | CollectionInfo info1; |
86 | CollectionInfo info2; |
87 | CollectionInfo info3; |
88 | info1.mutable_config()->set_collection_name("collection1" ); |
89 | info1.set_uuid("collection1-uuid" ); |
90 | info2.mutable_config()->set_collection_name("collection2" ); |
91 | info2.set_uuid("collection2-uuid" ); |
92 | info3.mutable_config()->set_collection_name("collection3" ); |
93 | info3.set_uuid("collection3-uuid" ); |
94 | std::vector<CollectionInfo> collection_infos{info1, info2, info3}; |
95 | std::vector<CollectionInfo> new_collections; |
96 | std::vector<std::string> old_collections; |
97 | std::vector<std::string> expired_collections; |
98 | collection_manager->classify_collections(collection_infos, &new_collections, |
99 | &old_collections, |
100 | &expired_collections); |
101 | EXPECT_EQ(3, new_collections.size()); |
102 | EXPECT_EQ(0, old_collections.size()); |
103 | EXPECT_EQ(0, expired_collections.size()); |
104 | EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(), |
105 | [](const CollectionInfo &info) { |
106 | return info.config().collection_name() == |
107 | "collection1" && |
108 | info.uuid() == "collection1-uuid" ; |
109 | }), |
110 | true); |
111 | EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(), |
112 | [](const CollectionInfo &info) { |
113 | return info.config().collection_name() == |
114 | "collection2" && |
115 | info.uuid() == "collection2-uuid" ; |
116 | }), |
117 | true); |
118 | EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(), |
119 | [](const CollectionInfo &info) { |
120 | return info.config().collection_name() == |
121 | "collection3" && |
122 | info.uuid() == "collection3-uuid" ; |
123 | }), |
124 | true); |
125 | collection_manager->create_collections(new_collections); |
126 | size_t count = 0; |
127 | auto created_collection = svc_.get_created_collections(); |
128 | auto collections_name = svc_.get_collections_name(); |
129 | while (count < 5) { // in order to sleep enough time |
130 | if (created_collection.size() == collections_name.size()) { |
131 | break; |
132 | } |
133 | std::this_thread::sleep_for(std::chrono::milliseconds(1000)); |
134 | created_collection = svc_.get_created_collections(); |
135 | collections_name = svc_.get_collections_name(); |
136 | count++; |
137 | } |
138 | EXPECT_EQ(created_collection.size(), |
139 | collection_manager->collections_.size()); |
140 | collection_manager->update_collections(old_collections); |
141 | collection_manager->drop_collections(expired_collections); |
142 | EXPECT_EQ(3, collection_manager->collections_.size()); |
143 | EXPECT_EQ(1, collection_manager->collections_.count("collection1-uuid" )); |
144 | EXPECT_EQ(1, collection_manager->collections_.count("collection2-uuid" )); |
145 | EXPECT_EQ(1, collection_manager->collections_.count("collection3-uuid" )); |
146 | |
147 | // Drop |
148 | collection_infos.clear(); |
149 | new_collections.clear(); |
150 | old_collections.clear(); |
151 | expired_collections.clear(); |
152 | collection_manager->classify_collections(collection_infos, &new_collections, |
153 | &old_collections, |
154 | &expired_collections); |
155 | EXPECT_EQ(0, new_collections.size()); |
156 | EXPECT_EQ(0, old_collections.size()); |
157 | EXPECT_EQ(3, expired_collections.size()); |
158 | EXPECT_EQ( |
159 | std::any_of( |
160 | expired_collections.begin(), expired_collections.end(), |
161 | [](const std::string &name) { return name == "collection1-uuid" ; }), |
162 | true); |
163 | EXPECT_EQ( |
164 | std::any_of( |
165 | expired_collections.begin(), expired_collections.end(), |
166 | [](const std::string &name) { return name == "collection2-uuid" ; }), |
167 | true); |
168 | EXPECT_EQ( |
169 | std::any_of( |
170 | expired_collections.begin(), expired_collections.end(), |
171 | [](const std::string &name) { return name == "collection3-uuid" ; }), |
172 | true); |
173 | collection_manager->create_collections(new_collections); |
174 | collection_manager->update_collections(old_collections); |
175 | collection_manager->drop_collections(expired_collections); |
176 | collection_manager->clean_invalid_collections(); |
177 | EXPECT_EQ(0, collection_manager->collections_.size()); |
178 | ret = collection_manager->stop(); |
179 | EXPECT_EQ(ret, 0); |
180 | } |
181 | ASSERT_EQ(0, server_.Stop(0)); |
182 | ASSERT_EQ(0, server_.Join()); |
183 | } |
184 | |