From efcc00f51cdeca73a3b4c1b5396e334836e728af Mon Sep 17 00:00:00 2001 From: sunby Date: Wed, 26 Jul 2023 17:59:03 +0800 Subject: [PATCH] Add some test utils macro Signed-off-by: sunby --- cpp/src/storage/test_util.h | 11 ----- cpp/test/space_test.cpp | 61 +++++++++---------------- cpp/{src/storage => test}/test_util.cpp | 2 +- cpp/test/test_util.h | 28 ++++++++++++ 4 files changed, 51 insertions(+), 51 deletions(-) delete mode 100644 cpp/src/storage/test_util.h rename cpp/{src/storage => test}/test_util.cpp (93%) create mode 100644 cpp/test/test_util.h diff --git a/cpp/src/storage/test_util.h b/cpp/src/storage/test_util.h deleted file mode 100644 index 74bc6445..00000000 --- a/cpp/src/storage/test_util.h +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -#include -#include -#include -#include "arrow/type.h" - -namespace milvus_storage { -std::shared_ptr CreateArrowSchema(std::vector field_names, - std::vector> field_types); -} // namespace milvus_storage \ No newline at end of file diff --git a/cpp/test/space_test.cpp b/cpp/test/space_test.cpp index ea2e5652..c81b2e12 100644 --- a/cpp/test/space_test.cpp +++ b/cpp/test/space_test.cpp @@ -8,71 +8,54 @@ #include #include #include +#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{"key1", "key2"}, - std::vector{"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(); 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{"key1", "key2"}, std::vector{"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(arrow_schema, schema_options); - auto sp_status = space_schema->Validate(); - ASSERT_TRUE(sp_status.ok()); + auto schema = std::make_shared(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 pk_array; - pk_builder.Finish(&pk_array); + ASSERT_STATUS_OK(pk_builder.Finish(&pk_array)); std::shared_ptr ts_array; - ts_builder.Finish(&ts_array); + ASSERT_STATUS_OK(ts_builder.Finish(&ts_array)); std::shared_ptr 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 b/cpp/test/test_util.cpp similarity index 93% rename from cpp/src/storage/test_util.cpp rename to cpp/test/test_util.cpp index a5d9d840..5dae780e 100644 --- a/cpp/src/storage/test_util.cpp +++ b/cpp/test/test_util.cpp @@ -8,4 +8,4 @@ std::shared_ptr CreateArrowSchema(std::vector field_ } return std::make_shared(fields); } -} // namespace milvus_storage \ No newline at end of file +} // namespace milvus_storage diff --git a/cpp/test/test_util.h b/cpp/test/test_util.h new file mode 100644 index 00000000..f9454307 --- /dev/null +++ b/cpp/test/test_util.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include +#include +#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 CreateArrowSchema(std::vector field_names, + std::vector> field_types); +} // namespace milvus_storage