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 CollectionManagerFilterTest : public ::testing::Test {
44 protected:
45 CollectionManagerFilterTest() {}
46
47 ~CollectionManagerFilterTest() {}
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(CollectionManagerFilterTest, 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