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 | |
17 | #include <ailego/encoding/json.h> |
18 | #include <ailego/utility/float_helper.h> |
19 | #include <ailego/utility/time_helper.h> |
20 | |
21 | #define private public |
22 | #include "admin/admin_proto_converter.h" |
23 | #undef private |
24 | #include <gmock/gmock.h> |
25 | #include <gtest/gtest.h> |
26 | |
27 | using namespace std; |
28 | using namespace testing; |
29 | |
30 | namespace proxima { |
31 | namespace be { |
32 | namespace admin { |
33 | namespace test { |
34 | |
35 | proto::CollectionConfig GetTestCollectionConfig() { |
36 | proto::CollectionConfig config; |
37 | config.set_collection_name("collection" ); |
38 | config.set_max_docs_per_segment(1000); |
39 | config.add_forward_column_names("f1" ); |
40 | config.add_forward_column_names("f2" ); |
41 | config.add_forward_column_names("f3" ); |
42 | auto *index = config.add_index_column_params(); |
43 | index->set_dimension(32); |
44 | index->set_column_name("index1" ); |
45 | index->set_data_type(proto::DT_VECTOR_FP16); |
46 | return config; |
47 | } |
48 | |
49 | proto::CollectionConfig GetTestCollectionConfigWithDbRepository() { |
50 | auto config = GetTestCollectionConfig(); |
51 | auto *repo = config.mutable_repository_config(); |
52 | repo->set_repository_name("test_repo" ); |
53 | repo->set_repository_type( |
54 | proto::CollectionConfig_RepositoryConfig_RepositoryType_RT_DATABASE); |
55 | auto *db = repo->mutable_database(); |
56 | db->set_password("password" ); |
57 | db->set_user("user" ); |
58 | db->set_connection_uri("url" ); |
59 | db->set_table_name("table" ); |
60 | return config; |
61 | } |
62 | |
63 | |
64 | TEST(AdminProtoConverterTest, PBToCollectionBase) { |
65 | auto config = GetTestCollectionConfig(); |
66 | meta::CollectionBase c; |
67 | ASSERT_EQ(AdminProtoConverter::PBToCollectionBase(config, &c), 0); |
68 | EXPECT_EQ(c.name(), "collection" ); |
69 | EXPECT_EQ(c.max_docs_per_segment(), 1000); |
70 | EXPECT_THAT(c.forward_columns(), |
71 | testing::ContainerEq(vector<string>{"f1" , "f2" , "f3" })); |
72 | EXPECT_FALSE(c.repository()); |
73 | auto &index_columns = c.index_columns(); |
74 | EXPECT_EQ(index_columns.size(), 1U); |
75 | auto &index = index_columns[0]; |
76 | EXPECT_EQ(index->name(), "index1" ); |
77 | EXPECT_EQ(index->dimension(), 32); |
78 | EXPECT_EQ(index->data_type(), DataTypes::VECTOR_FP16); |
79 | }; |
80 | |
81 | TEST(AdminProtoConverterTest, PBToCollectionBaseWithRepository) { |
82 | auto config = GetTestCollectionConfigWithDbRepository(); |
83 | meta::CollectionBase c; |
84 | ASSERT_EQ(AdminProtoConverter::PBToCollectionBase(config, &c), 0); |
85 | EXPECT_EQ(c.name(), "collection" ); |
86 | EXPECT_EQ(c.max_docs_per_segment(), 1000); |
87 | EXPECT_THAT(c.forward_columns(), |
88 | testing::ContainerEq(vector<string>{"f1" , "f2" , "f3" })); |
89 | auto &index_columns = c.index_columns(); |
90 | EXPECT_EQ(index_columns.size(), 1U); |
91 | auto &index = index_columns[0]; |
92 | EXPECT_EQ(index->name(), "index1" ); |
93 | EXPECT_EQ(index->dimension(), 32); |
94 | EXPECT_EQ(index->data_type(), DataTypes::VECTOR_FP16); |
95 | EXPECT_TRUE(c.repository()); |
96 | auto db_repo = |
97 | dynamic_pointer_cast<meta::DatabaseRepositoryMeta>(c.repository()); |
98 | EXPECT_TRUE(db_repo); |
99 | EXPECT_EQ(db_repo->password(), "password" ); |
100 | EXPECT_EQ(db_repo->user(), "user" ); |
101 | EXPECT_EQ(db_repo->table_name(), "table" ); |
102 | EXPECT_EQ(db_repo->connection(), "url" ); |
103 | EXPECT_EQ(db_repo->name(), "test_repo" ); |
104 | EXPECT_EQ(c.repository()->type(), meta::RepositoryTypes::DATABASE); |
105 | }; |
106 | |
107 | TEST(AdminProtoConverterTest, CollectionMetaToPB) { |
108 | auto config = GetTestCollectionConfig(); |
109 | meta::CollectionBase c; |
110 | ASSERT_EQ(AdminProtoConverter::PBToCollectionBase(config, &c), 0); |
111 | meta::CollectionMeta meta(c); |
112 | meta.set_status(meta::CollectionStatus::SERVING); |
113 | proto::CollectionInfo info; |
114 | AdminProtoConverter::CollectionMetaToPB(meta, &info); |
115 | EXPECT_EQ(info.status(), proto::CollectionInfo_CollectionStatus_CS_SERVING); |
116 | auto &conf = info.config(); |
117 | EXPECT_EQ(conf.max_docs_per_segment(), 1000); |
118 | EXPECT_EQ(conf.forward_column_names_size(), 3); |
119 | EXPECT_EQ(conf.forward_column_names(0), "f1" ); |
120 | EXPECT_EQ(conf.forward_column_names(1), "f2" ); |
121 | EXPECT_EQ(conf.forward_column_names(2), "f3" ); |
122 | EXPECT_EQ(conf.index_column_params_size(), 1); |
123 | auto &index = conf.index_column_params(0); |
124 | EXPECT_EQ(index.dimension(), 32); |
125 | EXPECT_EQ(index.data_type(), proto::DT_VECTOR_FP16); |
126 | EXPECT_EQ(index.column_name(), "index1" ); |
127 | EXPECT_EQ(conf.collection_name(), "collection" ); |
128 | EXPECT_FALSE(conf.has_repository_config()); |
129 | } |
130 | |
131 | TEST(AdminProtoConverterTest, CollectionMetaToPBWithRepository) { |
132 | auto config = GetTestCollectionConfigWithDbRepository(); |
133 | meta::CollectionBase c; |
134 | ASSERT_EQ(AdminProtoConverter::PBToCollectionBase(config, &c), 0); |
135 | meta::CollectionMeta meta(c); |
136 | meta.set_status(meta::CollectionStatus::SERVING); |
137 | proto::CollectionInfo info; |
138 | AdminProtoConverter::CollectionMetaToPB(meta, &info); |
139 | auto &conf = info.config(); |
140 | EXPECT_EQ(conf.max_docs_per_segment(), 1000); |
141 | EXPECT_EQ(conf.forward_column_names_size(), 3); |
142 | EXPECT_EQ(conf.forward_column_names(0), "f1" ); |
143 | EXPECT_EQ(conf.forward_column_names(1), "f2" ); |
144 | EXPECT_EQ(conf.forward_column_names(2), "f3" ); |
145 | EXPECT_EQ(conf.index_column_params_size(), 1); |
146 | auto &index = conf.index_column_params(0); |
147 | EXPECT_EQ(index.dimension(), 32); |
148 | EXPECT_EQ(index.data_type(), proto::DT_VECTOR_FP16); |
149 | EXPECT_EQ(index.column_name(), "index1" ); |
150 | EXPECT_EQ(conf.collection_name(), "collection" ); |
151 | EXPECT_TRUE(conf.has_repository_config()); |
152 | auto &repo_config = conf.repository_config(); |
153 | EXPECT_EQ(repo_config.repository_name(), "test_repo" ); |
154 | EXPECT_EQ( |
155 | repo_config.repository_type(), |
156 | proto::CollectionConfig_RepositoryConfig_RepositoryType_RT_DATABASE); |
157 | EXPECT_TRUE(repo_config.has_database()); |
158 | auto &db = repo_config.database(); |
159 | EXPECT_EQ(db.user(), "user" ); |
160 | EXPECT_EQ(db.password(), "password" ); |
161 | EXPECT_EQ(db.table_name(), "table" ); |
162 | EXPECT_EQ(db.connection_uri(), "url" ); |
163 | } |
164 | |
165 | TEST(AdminProtoConverterTest, CollectionStatsToPB) { |
166 | index::CollectionStats stats; |
167 | stats.collection_name = "collection" ; |
168 | stats.collection_path = "path" ; |
169 | stats.total_segment_count = 1000; |
170 | index::SegmentMeta sm; |
171 | sm.segment_id = 11; |
172 | sm.max_doc_id = 3000; |
173 | stats.segment_stats.push_back(index::SegmentStats(sm)); |
174 | stats.segment_stats.push_back(index::SegmentStats(sm)); |
175 | proto::CollectionStats pb_stats; |
176 | AdminProtoConverter::CollectionStatsToPB(stats, &pb_stats); |
177 | EXPECT_EQ(pb_stats.collection_name(), "collection" ); |
178 | EXPECT_EQ(pb_stats.collection_path(), "path" ); |
179 | EXPECT_EQ(pb_stats.total_segment_count(), 1000); |
180 | EXPECT_EQ(pb_stats.segment_stats_size(), 2); |
181 | EXPECT_EQ(pb_stats.segment_stats(0).segment_id(), 11); |
182 | EXPECT_EQ(pb_stats.segment_stats(0).max_doc_id(), 3000); |
183 | } |
184 | |
185 | } // namespace test |
186 | } // namespace admin |
187 | } // namespace be |
188 | } // namespace proxima |
189 | |