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 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 | |
63 | TEST_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 | |