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
27using namespace std;
28using namespace testing;
29
30namespace proxima {
31namespace be {
32namespace admin {
33namespace test {
34
35proto::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
49proto::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
64TEST(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
81TEST(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
107TEST(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
131TEST(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
165TEST(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