Skip to content

Commit

Permalink
fix: fix ShowPartitions error, milvus-io#246
Browse files Browse the repository at this point in the history
  • Loading branch information
matrixji committed Jan 14, 2024
1 parent 0127694 commit 3fb71a4
Show file tree
Hide file tree
Showing 8 changed files with 161 additions and 189 deletions.
6 changes: 5 additions & 1 deletion src/impl/MilvusClientImpl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -423,8 +423,12 @@ MilvusClientImpl::ShowPartitions(const std::string& collection_name, const std::
partitions_info.reserve(count);
}
for (int i = 0; i < count; ++i) {
auto inmemory_percentage = 0;
if (response.inmemory_percentages_size() > i) {
inmemory_percentage = response.inmemory_percentages(i);
}
partitions_info.emplace_back(response.partition_names(i), response.partitionids(i),
response.created_timestamps(i), response.inmemory_percentages(i));
response.created_timestamps(i), inmemory_percentage);
}
};

Expand Down
145 changes: 114 additions & 31 deletions test/st/MilvusServerTest.h
Original file line number Diff line number Diff line change
Expand Up @@ -29,36 +29,123 @@
namespace milvus {
namespace test {

inline void
waitMilvusServerReady(const PythonMilvusServer& server) {
int max_retry = 60, retry = 0;
bool has;

auto client = milvus::MilvusClient::Create();
auto param = server.TestClientParam();
client->Connect(*param);
auto status = client->HasCollection("no_such", has);

while (!status.IsOk() && retry++ < max_retry) {
std::this_thread::sleep_for(std::chrono::seconds{5});
client = milvus::MilvusClient::Create();
client->Connect(*param);
status = client->HasCollection("no_such", has);
std::cout << "Wait milvus start done, try: " << retry << ", status: " << status.Message() << std::endl;
}
std::cout << "Wait milvus start done, status: " << status.Message() << std::endl;
}

class MilvusServerTest : public ::testing::Test {
protected:
class MilvusServerClientPair {
private:
PythonMilvusServer server_{};
std::shared_ptr<milvus::MilvusClient> client_{nullptr};

public:
std::string collection_name{"Foo"};
std::string partition_name{"Bar"};

void
SetUp() override {
startServer() {
server_.Start();
client_ = milvus::MilvusClient::Create();
waitMilvusServerReady(server_);
waitMilvusServerReady();
}

PythonMilvusServer&
server() {
return server_;
}

milvus::MilvusClient&
client() {
return *client_;
}

void
waitMilvusServerReady() {
int max_retry = 60, retry = 0;
bool has;

auto client = milvus::MilvusClient::Create();
auto param = server_.TestClientParam();
client->Connect(*param);
auto status = client->HasCollection("no_such", has);

while (!status.IsOk() && retry++ < max_retry) {
std::this_thread::sleep_for(std::chrono::seconds{5});
client = milvus::MilvusClient::Create();
client->Connect(*param);
status = client->HasCollection("no_such", has);
std::cout << "Wait milvus start done, try: " << retry << ", status: " << status.Message() << std::endl;
}
std::cout << "Wait milvus start done, status: " << status.Message() << std::endl;
}

void
createBinaryCollectionAndPartitions() {
milvus::CollectionSchema collection_schema(collection_name);
collection_schema.AddField(milvus::FieldSchema("id", milvus::DataType::INT64, "id", true, true));
collection_schema.AddField(milvus::FieldSchema("age", milvus::DataType::INT16, "age"));
collection_schema.AddField(
milvus::FieldSchema("face", milvus::DataType::BINARY_VECTOR, "face signature").WithDimension(32));

auto status = client_->CreateCollection(collection_schema);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());

status = client_->CreatePartition(collection_name, partition_name);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());
}

void
createCollectionAndPartitions(bool create_flat_index) {
milvus::CollectionSchema collection_schema(collection_name);
collection_schema.AddField(milvus::FieldSchema("id", milvus::DataType::INT64, "id", true, true));
collection_schema.AddField(milvus::FieldSchema("age", milvus::DataType::INT16, "age"));
collection_schema.AddField(milvus::FieldSchema("name", milvus::DataType::VARCHAR, "name").WithMaxLength(64));
collection_schema.AddField(
milvus::FieldSchema("face", milvus::DataType::FLOAT_VECTOR, "face signature").WithDimension(4));

auto status = client_->CreateCollection(collection_schema);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());

if (create_flat_index) {
milvus::IndexDesc index_desc("face", "", milvus::IndexType::FLAT, milvus::MetricType::L2, 0);
status = client_->CreateIndex(collection_name, index_desc);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());
}

status = client_->CreatePartition(collection_name, partition_name);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());
}

milvus::DmlResults
insertRecords(const std::vector<milvus::FieldDataPtr>& fields) {
milvus::DmlResults dml_results;
auto status = client_->Insert(collection_name, partition_name, fields, dml_results);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());
EXPECT_EQ(dml_results.IdArray().IntIDArray().size(), fields.front()->Count());
return dml_results;
}

void
loadCollection() {
auto status = client_->LoadCollection(collection_name);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());
}

void
dropCollection() {
auto status = client_->DropCollection(collection_name);
EXPECT_TRUE(status.IsOk());
}
};

class MilvusServerTest : public ::testing::Test, public MilvusServerClientPair {
protected:
void
SetUp() override {
startServer();
}

void
Expand All @@ -67,21 +154,17 @@ class MilvusServerTest : public ::testing::Test {
};

template <typename T>
class MilvusServerTestWithParam : public ::testing::TestWithParam<T> {
class MilvusServerTestWithParam : public ::testing::TestWithParam<T>, public MilvusServerClientPair {
protected:
PythonMilvusServer server_{};
std::shared_ptr<milvus::MilvusClient> client_{nullptr};

void
SetUp() override {
server_.Start();
client_ = milvus::MilvusClient::Create();
waitMilvusServerReady(server_);
startServer();
}

void
TearDown() override {
}
};

} // namespace test
} // namespace milvus
23 changes: 14 additions & 9 deletions test/st/TestCollection.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ using MilvusServerTestCollection = MilvusServerTestWithParam<bool>;

TEST_P(MilvusServerTestCollection, CreateAndDeleteCollection) {
auto using_string_primary_key = GetParam();
milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()};
client_->Connect(connect_param);
milvus::ConnectParam connect_param{"127.0.0.1", server().ListenPort()};
client().Connect(connect_param);

milvus::CollectionSchema collection_schema("Foo");
if (using_string_primary_key) {
Expand All @@ -38,38 +38,43 @@ TEST_P(MilvusServerTestCollection, CreateAndDeleteCollection) {
collection_schema.AddField(
milvus::FieldSchema("face", milvus::DataType::FLOAT_VECTOR, "face signature").WithDimension(1024));

auto status = client_->CreateCollection(collection_schema);
auto status = client().CreateCollection(collection_schema);
EXPECT_EQ(status.Message(), "OK");
EXPECT_TRUE(status.IsOk());

// create index needed after 2.2.0
milvus::IndexDesc index_desc("face", "", milvus::IndexType::FLAT, milvus::MetricType::L2, 0);
status = client_->CreateIndex("Foo", index_desc);
status = client().CreateIndex("Foo", index_desc);
EXPECT_TRUE(status.IsOk());

// test for https://github.com/milvus-io/milvus-sdk-cpp/issues/188
std::vector<std::string> names;
std::vector<milvus::CollectionInfo> collection_infos;
status = client_->ShowCollections(names, collection_infos);
status = client().ShowCollections(names, collection_infos);
EXPECT_TRUE(status.IsOk());
EXPECT_EQ(collection_infos.size(), 1);
EXPECT_EQ(collection_infos.front().MemoryPercentage(), 0);
EXPECT_EQ(collection_infos.front().Name(), "Foo");

// test for https://github.com/milvus-io/milvus-sdk-cpp/issues/246
milvus::PartitionsInfo partitionsInfo{};
status = client().ShowPartitions(collection_name, std::vector<std::string>{}, partitionsInfo);
EXPECT_TRUE(status.IsOk());

names.emplace_back("Foo");
collection_infos.clear();
status = client_->LoadCollection("Foo");
status = client().LoadCollection("Foo");
EXPECT_TRUE(status.IsOk());

status = client_->ShowCollections(names, collection_infos);
status = client().ShowCollections(names, collection_infos);
EXPECT_TRUE(status.IsOk());
EXPECT_EQ(collection_infos.size(), 1);
EXPECT_EQ(collection_infos.front().MemoryPercentage(), 100);

status = client_->RenameCollection("Foo", "Bar");
status = client().RenameCollection("Foo", "Bar");
EXPECT_TRUE(status.IsOk());

status = client_->DropCollection("Bar");
status = client().DropCollection("Bar");
EXPECT_TRUE(status.IsOk());
}

Expand Down
16 changes: 8 additions & 8 deletions test/st/TestConnectWithTls.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -73,12 +73,12 @@ class MilvusServerTestWithTlsMode : public MilvusServerTest {
std::array<char, 256> path{};
getcwd(path.data(), path.size());
std::string pwd = path.data();
server_.SetTls(Mode, pwd + "/certs/server.crt", pwd + "/certs/server.key", pwd + "/certs/ca.crt");
server().SetTls(Mode, pwd + "/certs/server.crt", pwd + "/certs/server.key", pwd + "/certs/ca.crt");

MilvusServerTest::SetUp();

ssl_client_ = milvus::MilvusClient::Create();
auto param = server_.TestClientParam();
auto param = server().TestClientParam();
ssl_client_->Connect(*param);
}

Expand All @@ -95,25 +95,25 @@ class DISABLED_MilvusServerTestWithTlsMode2 : public MilvusServerTestWithTlsMode

TEST_F(MilvusServerTestWithTlsMode1, GenericTest) {
bool has;
auto status = ssl_client_->HasCollection("nosuchcollection", has);
auto status = ssl_client_->HasCollection("no_such", has);
EXPECT_TRUE(status.IsOk());
status = client_->HasCollection("nosuchcollection", has);
status = client().HasCollection("no_such", has);
EXPECT_FALSE(status.IsOk());
EXPECT_EQ(status.Code(), milvus::StatusCode::NOT_CONNECTED);
}

TEST_F(DISABLED_MilvusServerTestWithTlsMode2, GenericTest) {
bool has;
auto status = ssl_client_->HasCollection("nosuchcollection", has);
auto status = ssl_client_->HasCollection("no_such", has);
EXPECT_TRUE(status.IsOk());
status = client_->HasCollection("nosuchcollection", has);
status = client().HasCollection("no_such", has);
EXPECT_FALSE(status.IsOk());
EXPECT_EQ(status.Code(), milvus::StatusCode::NOT_CONNECTED);

milvus::ConnectParam param{"127.0.0.1", 300};
param.EnableTls();
client_->Connect(param);
status = client_->HasCollection("nosuchcollection", has);
client().Connect(param);
status = client().HasCollection("no_such", has);
EXPECT_FALSE(status.IsOk());
EXPECT_EQ(status.Code(), milvus::StatusCode::NOT_CONNECTED);
}
8 changes: 4 additions & 4 deletions test/st/TestConnectWithUser.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,10 @@ class MilvusServerTestWithAuth : public MilvusServerTest {

void
SetUp() override {
server_.SetAuthorizationEnabled(true);
server().SetAuthorizationEnabled(true);
MilvusServerTest::SetUp();
root_client_ = milvus::MilvusClient::Create();
root_client_->Connect({"127.0.0.1", server_.ListenPort(), "root", "Milvus"});
root_client_->Connect({"127.0.0.1", server().ListenPort(), "root", "Milvus"});
}

void
Expand All @@ -46,8 +46,8 @@ TEST_F(MilvusServerTestWithAuth, GenericTest) {
bool has;
auto status = root_client_->HasCollection("nosuchcollection", has);
EXPECT_TRUE(status.IsOk());
// orig client with out user/pass
status = client_->HasCollection("nosuchcollection", has);
// orig client without user/pass
status = client().HasCollection("nosuchcollection", has);
EXPECT_FALSE(status.IsOk());
EXPECT_EQ(status.Code(), milvus::StatusCode::NOT_CONNECTED);
}
6 changes: 3 additions & 3 deletions test/st/TestGeneric.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,11 @@ using milvus::test::MilvusServerTest;
class MilvusServerTestGeneric : public MilvusServerTest {};

TEST_F(MilvusServerTestGeneric, GetVersion) {
milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()};
client_->Connect(connect_param);
milvus::ConnectParam connect_param{"127.0.0.1", server().ListenPort()};
client().Connect(connect_param);

std::string version;
auto status = client_->GetVersion(version);
auto status = client().GetVersion(version);
EXPECT_TRUE(status.IsOk());
EXPECT_THAT(version, testing::StartsWith("v2."));
}
Loading

0 comments on commit 3fb71a4

Please sign in to comment.