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 CollectionManagerCreateTest : public ::testing::Test {
44 protected:
45 CollectionManagerCreateTest() {}
46
47 ~CollectionManagerCreateTest() {}
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
63TEST_F(CollectionManagerCreateTest, TestCreate) {
64 brpc::Server server_;
65 MockProximaServiceImpl svc_;
66 brpc::ServerOptions options;
67 ASSERT_EQ(0, server_.AddService(&svc_, brpc::SERVER_DOESNT_OWN_SERVICE));
68 ASSERT_EQ(0, server_.Start(PORT, &options));
69 {
70 MockCollectionCreatorPtr collection_creator =
71 std::make_shared<MockCollectionCreator>();
72 EXPECT_CALL(*collection_creator, create(_))
73 .WillRepeatedly(Invoke([](const CollectionInfo &info) -> CollectionPtr {
74 std::string collection_name = info.config().collection_name();
75 MysqlHandlerPtr mysql_handler =
76 std::make_shared<MysqlHandler>(info.config());
77 return std::make_shared<FakeMysqlCollection>(info.config(),
78 mysql_handler);
79 }));
80
81 CollectionManagerPtr collection_manager =
82 std::make_shared<CollectionManager>(collection_creator);
83 int ret = collection_manager->init();
84 ASSERT_EQ(0, ret);
85
86 CollectionInfo info1;
87 CollectionInfo info2;
88 CollectionInfo info3;
89 info1.mutable_config()->set_collection_name("collection1");
90 info1.set_uuid("collection1-uuid");
91 info2.mutable_config()->set_collection_name("collection2");
92 info2.set_uuid("collection2-uuid");
93 info3.mutable_config()->set_collection_name("collection3");
94 info3.set_uuid("collection3-uuid");
95 std::vector<CollectionInfo> collection_infos{info1, info2, info3};
96 std::vector<CollectionInfo> new_collections;
97 std::vector<std::string> old_collections;
98 std::vector<std::string> expired_collections;
99 collection_manager->classify_collections(collection_infos, &new_collections,
100 &old_collections,
101 &expired_collections);
102 EXPECT_EQ(3, new_collections.size());
103 EXPECT_EQ(0, old_collections.size());
104 EXPECT_EQ(0, expired_collections.size());
105 EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(),
106 [](const CollectionInfo &info) {
107 return info.config().collection_name() ==
108 "collection1" &&
109 info.uuid() == "collection1-uuid";
110 }),
111 true);
112 EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(),
113 [](const CollectionInfo &info) {
114 return info.config().collection_name() ==
115 "collection2" &&
116 info.uuid() == "collection2-uuid";
117 }),
118 true);
119 EXPECT_EQ(std::any_of(new_collections.begin(), new_collections.end(),
120 [](const CollectionInfo &info) {
121 return info.config().collection_name() ==
122 "collection3" &&
123 info.uuid() == "collection3-uuid";
124 }),
125 true);
126
127 collection_manager->create_collections(new_collections);
128 size_t count = 0;
129 auto created_collection = svc_.get_created_collections();
130 auto collections_name = svc_.get_collections_name();
131 while (count < 5) { // in order to sleep enough time
132 if (created_collection.size() == collections_name.size()) {
133 break;
134 }
135 std::this_thread::sleep_for(std::chrono::milliseconds(1000));
136 created_collection = svc_.get_created_collections();
137 collections_name = svc_.get_collections_name();
138 count++;
139 }
140 EXPECT_EQ(created_collection.size(),
141 collection_manager->collections_.size());
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 ret = collection_manager->stop();
147 EXPECT_EQ(ret, 0);
148 }
149 ASSERT_EQ(0, server_.Stop(0));
150 ASSERT_EQ(0, server_.Join());
151}
152