Add some test utils macro (#31)

Signed-off-by: sunby <[email protected]>
milvus-io · Jul 27, 2023 · a7f850b · a7f850b
1 parent 941a649
commit a7f850b
Show file tree

Hide file tree

Showing 8 changed files with 63 additions and 62 deletions.
diff --git a/cpp/src/common/utils.cpp b/cpp/src/common/utils.cpp
@@ -39,7 +39,9 @@ Result<std::unique_ptr<schema_proto::Field>> ToProtobufField(const arrow::Field*
   auto proto_field = std::make_unique<schema_proto::Field>();
   proto_field->set_name(field->name());
   proto_field->set_nullable(field->nullable());
-  proto_field->set_allocated_metadata(ToProtobufMetadata(field->metadata().get()).release());
+  if (field->metadata() != nullptr) {
+    proto_field->set_allocated_metadata(ToProtobufMetadata(field->metadata().get()).release());
+  }
   ASSIGN_OR_RETURN_NOT_OK(auto data_type, ToProtobufDataType(field->type().get()));
 
   proto_field->set_allocated_data_type(data_type.release());
@@ -113,11 +115,13 @@ Result<std::unique_ptr<schema_proto::ArrowSchema>> ToProtobufSchema(const arrow:
   proto_schema->set_endianness(schema->endianness() == arrow::Endianness::Little ? schema_proto::Endianness::Little
                                                                                  : schema_proto::Endianness::Big);
 
-  for (const auto& key : schema->metadata()->keys()) {
-    proto_schema->mutable_metadata()->add_keys(key);
-  }
-  for (const auto& value : schema->metadata()->values()) {
-    proto_schema->mutable_metadata()->add_values(value);
+  if (schema->metadata() != nullptr) {
+    for (const auto& key : schema->metadata()->keys()) {
+      proto_schema->mutable_metadata()->add_keys(key);
+    }
+    for (const auto& value : schema->metadata()->values()) {
+      proto_schema->mutable_metadata()->add_values(value);
+    }
   }
   return proto_schema;
 }

diff --git a/cpp/src/storage/test_util.h b/cpp/src/storage/test_util.h
diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt
@@ -6,6 +6,7 @@ add_executable(
   schema_test.cpp
   manifest_test.cpp
   space_test.cpp
+  test_util.cpp
 )
 
 target_link_libraries(

diff --git a/cpp/test/options_test.cpp b/cpp/test/options_test.cpp
@@ -2,7 +2,7 @@
 #include "gtest/gtest.h"
 #include "storage/options.h"
 #include "arrow/type.h"
-#include "storage/test_util.h"
+#include "test_util.h"
 #include "google/protobuf/util/message_differencer.h"
 
 namespace milvus_storage {

diff --git a/cpp/test/schema_test.cpp b/cpp/test/schema_test.cpp
@@ -1,7 +1,7 @@
 #include "storage/schema.h"
 #include "gtest/gtest.h"
 #include "storage/options.h"
-#include "storage/test_util.h"
+#include "test_util.h"
 #include <arrow/util/key_value_metadata.h>
 
 namespace milvus_storage {

diff --git a/cpp/test/space_test.cpp b/cpp/test/space_test.cpp
@@ -1,4 +1,5 @@
 #include "storage/space.h"
+#include "common/log.h"
 #include "filter/constant_filter.h"
 
 #include <memory>
@@ -8,70 +9,48 @@
 #include <arrow/array/builder_primitive.h>
 #include <arrow/util/key_value_metadata.h>
 #include <gtest/gtest.h>
+#include "test_util.h"
 
 namespace milvus_storage {
 /**
  * @brief Test Space::Write
  *
  */
-TEST(SpaceTest, SpaceWriteTest) {
-  arrow::SchemaBuilder schema_builder;
-  auto metadata = arrow::KeyValueMetadata::Make(std::vector<std::string>{"key1", "key2"},
-                                                std::vector<std::string>{"value1", "value2"});
-
-  auto pk_field = arrow::field("pk_field", arrow::int64(), /*nullable=*/false, metadata);
-
-  auto ts_field = arrow::field("ts_field", arrow::int64(), /*nullable=*/false, metadata);
-
-  auto vec_field = arrow::field("vec_field", arrow::fixed_size_binary(10), /*nullable=*/false, metadata);
-  auto status = schema_builder.AddField(pk_field);
-  ASSERT_TRUE(status.ok());
-  status = schema_builder.AddField(ts_field);
-  ASSERT_TRUE(status.ok());
-  status = schema_builder.AddField(vec_field);
-  ASSERT_TRUE(status.ok());
+TEST(SpaceTest, SpaceWriteReadTest) {
+  auto arrow_schema = CreateArrowSchema({"pk_field", "ts_field", "vec_field"},
+                                        {arrow::int64(), arrow::int64(), arrow::fixed_size_binary(10)});
 
   auto schema_options = std::make_shared<SchemaOptions>();
   schema_options->primary_column = "pk_field";
   schema_options->version_column = "ts_field";
   schema_options->vector_column = "vec_field";
 
-  auto schema_metadata =
-      arrow::KeyValueMetadata(std::vector<std::string>{"key1", "key2"}, std::vector<std::string>{"value1", "value2"});
-  auto metadata_status = schema_builder.AddMetadata(schema_metadata);
-  ASSERT_TRUE(metadata_status.ok());
-
-  auto arrow_schema = schema_builder.Finish().ValueOrDie();
-
-  auto space_schema = std::make_shared<Schema>(arrow_schema, schema_options);
-  auto sp_status = space_schema->Validate();
-  ASSERT_TRUE(sp_status.ok());
+  auto schema = std::make_shared<Schema>(arrow_schema, schema_options);
+  ASSERT_STATUS_OK(schema->Validate());
 
   auto uri = "file:///tmp/";
-  auto res = Space::Open(uri, Options{space_schema, -1});
-  ASSERT_TRUE(res.ok());
-  auto space = std::move(res.value());
+  ASSERT_AND_ASSIGN(auto space, Space::Open(uri, Options{schema, -1}));
+  LogOut("xx");
 
-  // Create RecordBatch
   arrow::Int64Builder pk_builder;
-  pk_builder.Append(1);
-  pk_builder.Append(2);
-  pk_builder.Append(3);
+  ASSERT_STATUS_OK(pk_builder.Append(1));
+  ASSERT_STATUS_OK(pk_builder.Append(2));
+  ASSERT_STATUS_OK(pk_builder.Append(3));
   arrow::Int64Builder ts_builder;
-  ts_builder.Append(1);
-  ts_builder.Append(2);
-  ts_builder.Append(3);
+  ASSERT_STATUS_OK(ts_builder.Append(1));
+  ASSERT_STATUS_OK(ts_builder.Append(2));
+  ASSERT_STATUS_OK(ts_builder.Append(3));
   arrow::FixedSizeBinaryBuilder vec_builder(arrow::fixed_size_binary(10));
-  vec_builder.Append("1234567890");
-  vec_builder.Append("1234567890");
-  vec_builder.Append("1234567890");
+  ASSERT_STATUS_OK(vec_builder.Append("1234567890"));
+  ASSERT_STATUS_OK(vec_builder.Append("1234567890"));
+  ASSERT_STATUS_OK(vec_builder.Append("1234567890"));
 
   std::shared_ptr<arrow::Array> pk_array;
-  pk_builder.Finish(&pk_array);
+  ASSERT_STATUS_OK(pk_builder.Finish(&pk_array));
   std::shared_ptr<arrow::Array> ts_array;
-  ts_builder.Finish(&ts_array);
+  ASSERT_STATUS_OK(ts_builder.Finish(&ts_array));
   std::shared_ptr<arrow::Array> vec_array;
-  vec_builder.Finish(&vec_array);
+  ASSERT_STATUS_OK(vec_builder.Finish(&vec_array));
 
   auto rec_batch = arrow::RecordBatch::Make(arrow_schema, 3, {pk_array, ts_array, vec_array});
   auto reader = arrow::RecordBatchReader::Make({rec_batch}, arrow_schema).ValueOrDie();

diff --git a/cpp/src/storage/test_util.cpp → cpp/test/test_util.cpp b/cpp/src/storage/test_util.cpp → cpp/test/test_util.cpp
@@ -8,4 +8,4 @@ std::shared_ptr<arrow::Schema> CreateArrowSchema(std::vector<std::string> field_
   }
   return std::make_shared<arrow::Schema>(fields);
 }
-}  // namespace milvus_storage
+}  // namespace milvus_storage
diff --git a/cpp/test/test_util.h b/cpp/test/test_util.h
@@ -0,0 +1,28 @@
+#pragma once
+
+#include <arrow/type_fwd.h>
+#include <vector>
+#include <string>
+#include "arrow/type.h"
+
+namespace milvus_storage {
+#define ASSERT_STATUS_OK(status) \
+  do {                           \
+    ASSERT_TRUE((status).ok());  \
+  } while (false)
+
+#define ASSERT_STATUS_NOT_OK(status) \
+  do {                               \
+    ASSERT_FALSE((status).ok());     \
+  } while (false)
+
+#define ASSERT_AND_ASSIGN_IMPL(status_name, lhs, rexpr) \
+  auto status_name = (rexpr);                           \
+  ASSERT_STATUS_OK(status_name.status());               \
+  lhs = std::move(status_name).value();
+
+#define ASSERT_AND_ASSIGN(lhs, rexpr) ASSERT_AND_ASSIGN_IMPL(CONCAT(_tmp_value, __COUNTER__), lhs, rexpr);
+
+std::shared_ptr<arrow::Schema> CreateArrowSchema(std::vector<std::string> field_names,
+                                                 std::vector<std::shared_ptr<arrow::DataType>> field_types);
+}  // namespace milvus_storage