Add some test utils macro

Signed-off-by: sunby <[email protected]>
milvus-io · Jul 26, 2023 · efcc00f · efcc00f
1 parent ac9653f
commit efcc00f
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 51 deletions.
diff --git a/cpp/src/storage/test_util.h b/cpp/src/storage/test_util.h
diff --git a/cpp/test/space_test.cpp b/cpp/test/space_test.cpp
@@ -8,71 +8,54 @@
 #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/";
+<<<<<<< Updated upstream
   auto res = Space::Open(uri, Options{space_schema, -1});
   std::cout << res.status().ToString() << std::endl;
   ASSERT_TRUE(res.ok());
   auto space = std::move(res.value());
+=======
+  ASSERT_AND_ASSIGN(auto space, Space::Open(uri, Options{schema, -1}));
+>>>>>>> Stashed changes
 
-  // 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