// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements. See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership. The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License. You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied. See the License for the

// specific language governing permissions and limitations

// under the License.

#include "meta-store/versioned_value.h"

#include <bthread/bthread.h>

#include <fmt/format.h>

#include <foundationdb/fdb_c_options.g.h>

#include <gen_cpp/cloud.pb.h>

#include <gen_cpp/olap_common.pb.h>

#include <gen_cpp/olap_file.pb.h>

#include <gtest/gtest.h>

#include <iostream>

#include <memory>

#include <string>

#include <string_view>

#include "common/config.h"

#include "common/util.h"

#include "meta-store/codec.h"

#include "meta-store/mem_txn_kv.h"

#include "meta-store/txn_kv.h"

#include "meta-store/txn_kv_error.h"

#include "meta-store/versionstamp.h"

using namespace doris::cloud;

int main(int argc, char** argv) {

config::init(nullptr, true);

::testing::InitGoogleTest(&argc, argv);

return RUN_ALL_TESTS();

}

static size_t count_range(TxnKv* txn_kv, std::string_view begin = "",

std::string_view end = "\xFF") {

std::unique_ptr<Transaction> txn;

EXPECT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

if (!txn) {

return 0; // Failed to create transaction

}

FullRangeGetOptions opts;

opts.txn = txn.get();

auto iter = txn_kv->full_range_get(std::string(begin), std::string(end), std::move(opts));

size_t total = 0;

for (auto&& kvp = iter->next(); kvp.has_value(); kvp = iter->next()) {

total += 1;

}

EXPECT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

return total;

}

static std::string dump_range(TxnKv* txn_kv, std::string_view begin = "",

std::string_view end = "\xFF") {

std::unique_ptr<Transaction> txn;

if (txn_kv->create_txn(&txn) != TxnErrorCode::TXN_OK) {

return "Failed to create dump range transaction";

}

FullRangeGetOptions opts;

opts.txn = txn.get();

auto iter = txn_kv->full_range_get(std::string(begin), std::string(end), std::move(opts));

std::string buffer;

for (auto&& kv = iter->next(); kv.has_value(); kv = iter->next()) {

buffer += fmt::format("Key: {}, Value: {}\n", hex(kv->first), hex(kv->second));

}

EXPECT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

return buffer;

}

TEST(VersionedValueTest, Remove) {

auto txn_kv = std::make_shared<MemTxnKv>();

std::string key_prefix = "remove_key_";

// Put 10 versions

for (int i = 0; i < 10; ++i) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionPB version_pb;

version_pb.set_version(i * 10);

Versionstamp versionstamp(i * 10, 0);

versioned_put(txn.get(), key_prefix, versionstamp, version_pb.SerializeAsString());

std::string key(key_prefix);

encode_versionstamp(versionstamp, &key);

std::cout << "Put key: " << hex(key) << ", version: " << versionstamp.version()

<< std::endl;

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_EQ(count_range(txn_kv.get()), 10) << dump_range(txn_kv.get());

{

// Remove the latest version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

Versionstamp last_version(90, 0);

versioned_remove(txn.get(), key_prefix, last_version);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_EQ(count_range(txn_kv.get()), 9) << dump_range(txn_kv.get());

{

// Remove a specific version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

versioned_remove(txn.get(), key_prefix, Versionstamp(50, 0));

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_EQ(count_range(txn_kv.get()), 8) << dump_range(txn_kv.get());

{

// Verify the removed key is not found

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

Versionstamp snapshot(51, 0);

Versionstamp versionstamp;

std::string value;

ASSERT_EQ(versioned_get(txn.get(), key_prefix, snapshot, &versionstamp, &value),

TxnErrorCode::TXN_OK)

<< dump_range(txn_kv.get());

VersionPB version_pb;

ASSERT_TRUE(version_pb.ParseFromString(value));

ASSERT_EQ(versionstamp.version(), 40)

<< dump_range(txn_kv.get()) << versionstamp.to_string();

}

}

TEST(VersionedValueTest, Get) {

auto txn_kv = std::make_shared<MemTxnKv>();

std::string key_prefix = "get_key_";

// Put 9 versions

for (int i = 1; i < 10; ++i) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionPB version_pb;

version_pb.set_version(i * 10);

Versionstamp versionstamp(i * 10, 0);

versioned_put(txn.get(), key_prefix, versionstamp, version_pb.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

{

// Get with the latest version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

Versionstamp versionstamp;

std::string value;

ASSERT_EQ(versioned_get(txn.get(), key_prefix, Versionstamp::max(), &versionstamp, &value),

TxnErrorCode::TXN_OK);

ASSERT_EQ(versionstamp.version(), 90);

VersionPB version_pb;

ASSERT_TRUE(version_pb.ParseFromString(value));

ASSERT_EQ(version_pb.version(), 90);

}

{

// Get with a specific version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

std::string value;

Versionstamp versionstamp;

ASSERT_EQ(versioned_get(txn.get(), key_prefix, Versionstamp(51, 0), &versionstamp, &value),

TxnErrorCode::TXN_OK);

ASSERT_EQ(versionstamp.version(), 50);

VersionPB version_pb;

ASSERT_TRUE(version_pb.ParseFromString(value));

ASSERT_EQ(version_pb.version(), 50);

}

{

// Get with a non-existing version

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

std::string value;

Versionstamp versionstamp;

ASSERT_EQ(versioned_get(txn.get(), key_prefix, Versionstamp(5, 0), &versionstamp, &value),

TxnErrorCode::TXN_KEY_NOT_FOUND);

}

}

TEST(VersionedValueTest, GetRange) {

auto txn_kv = std::make_shared<MemTxnKv>();

std::string key_prefix = "get_range_key_";

{

// Put a series of versioned documents

std::vector<std::pair<int, Versionstamp>> versions = {

{10, Versionstamp(10, 0)}, {20, Versionstamp(20, 0)}, {30, Versionstamp(30, 0)},

{40, Versionstamp(40, 0)}, {10, Versionstamp(40, 0)}, {50, Versionstamp(50, 0)},

{30, Versionstamp(60, 0)}, {20, Versionstamp(70, 0)}, {10, Versionstamp(80, 0)},

{60, Versionstamp(90, 0)}, {70, Versionstamp(90, 0)},

};

for (const auto& [suffix, versionstamp] : versions) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionPB version_pb;

version_pb.set_version(suffix);

std::string key = fmt::format("{}{:02}", key_prefix, suffix);

versioned_put(txn.get(), key, versionstamp, version_pb.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

}

struct TestCase {

Versionstamp snapshot_version;

std::string begin_key;

std::string end_key;

std::vector<std::pair<int, Versionstamp>> expected_results;

};

std::vector<TestCase> test_cases = {

// Test case 1: Get all versions with max snapshot version

// Results should be in reverse order by key (70, 60, 50, 40, 30, 20, 10)

{

Versionstamp::max(),

key_prefix,

key_prefix + '\xFF' + '\xFF',

{{70, Versionstamp(90, 0)},

{60, Versionstamp(90, 0)},

{50, Versionstamp(50, 0)},

{40, Versionstamp(40, 0)},

{30, Versionstamp(60, 0)},

{20, Versionstamp(70, 0)},

{10, Versionstamp(80, 0)}},

},

// Test case 2: Get versions with a specific snapshot version (50, 0)

// Only keys with version < (50, 0) should be returned: 10->40, 20->20, 30->30, 40->40

// Results should be in reverse order by key (40, 30, 20, 10)

{

Versionstamp(50, 0),

key_prefix + "10",

key_prefix + "70",

{{40, Versionstamp(40, 0)},

{30, Versionstamp(30, 0)},

{20, Versionstamp(20, 0)},

{10, Versionstamp(40, 0)}},

},

// Test case 3: Get versions with a non-existing snapshot version

{Versionstamp(5, 0), key_prefix + "10", key_prefix + "70", {}},

// Test case 4: Get versions with a specific range [20, 60)

// Results should be in reverse order by key

{

Versionstamp::max(),

key_prefix + "20",

key_prefix + "60",

{{50, Versionstamp(50, 0)},

{40, Versionstamp(40, 0)},

{30, Versionstamp(60, 0)},

{20, Versionstamp(70, 0)}},

},

// Test case 5: Empty range - begin_key >= end_key

{

Versionstamp::max(),

key_prefix + "50",

key_prefix + "30",

{},

},

// Test case 6: Single key range [30, 31)

{

Versionstamp::max(),

key_prefix + "30",

key_prefix + "31",

{{30, Versionstamp(60, 0)}},

},

// Test case 7: Range with no matching keys [25, 30)

{

Versionstamp::max(),

key_prefix + "25",

key_prefix + "30",

{},

},

// Test case 8: Exact boundary - key 20 included, key 60 excluded

{

Versionstamp::max(),

key_prefix + "20",

key_prefix + "60",

{{50, Versionstamp(50, 0)},

{40, Versionstamp(40, 0)},

{30, Versionstamp(60, 0)},

{20, Versionstamp(70, 0)}},

},

// Test case 9: Version boundary test - snapshot at (45, 0)

// Should include versions < 45: 10->40, 20->20, 30->30, 40->40

{

Versionstamp(45, 0),

key_prefix,

key_prefix + '\xFF' + '\xFF',

{{40, Versionstamp(40, 0)},

{30, Versionstamp(30, 0)},

{20, Versionstamp(20, 0)},

{10, Versionstamp(40, 0)}},

},

// Test case 10: Test with exact version boundary (40, 0)

// Should include versions < 40: 10->10, 20->20, 30->30

{

Versionstamp(40, 0),

key_prefix,

key_prefix + '\xFF' + '\xFF',

{{30, Versionstamp(30, 0)},

{20, Versionstamp(20, 0)},

{10, Versionstamp(10, 0)}},

},

};

size_t case_index = 0;

for (auto&& tc : test_cases) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

// Get range with range [begin, end)

VersionedRangeGetOptions opts;

opts.snapshot_version = tc.snapshot_version;

opts.begin_key_selector = RangeKeySelector::FIRST_GREATER_OR_EQUAL;

opts.end_key_selector = RangeKeySelector::FIRST_GREATER_OR_EQUAL;

auto iter = versioned_get_range(txn.get(), tc.begin_key, tc.end_key, opts);

size_t count = 0;

for (auto&& kvp = iter->next(); kvp.has_value(); kvp = iter->next()) {

auto [key, version, _value] = kvp.value();

key.remove_prefix(key_prefix.size());

int suffix = std::stoi(std::string(key));

ASSERT_EQ(suffix, tc.expected_results[count].first) << " count=" << count;

ASSERT_EQ(version, tc.expected_results[count].second)

<< " count=" << count << ", key=" << hex(key)

<< ", version=" << version.version()

<< ", expected=" << tc.expected_results[count].second.version();

count += 1;

}

ASSERT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

ASSERT_EQ(count, tc.expected_results.size());

case_index += 1;

std::cout << "Test case " << case_index << " passed." << std::endl;

}

// Iterate cases via peek method

case_index = 0;

for (auto&& tc : test_cases) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

// Get range with range [begin, end)

VersionedRangeGetOptions opts;

opts.snapshot_version = tc.snapshot_version;

opts.begin_key_selector = RangeKeySelector::FIRST_GREATER_OR_EQUAL;

opts.end_key_selector = RangeKeySelector::FIRST_GREATER_OR_EQUAL;

auto iter = versioned_get_range(txn.get(), tc.begin_key, tc.end_key, opts);

size_t count = 0;

for (auto&& kvp = iter->peek(); kvp.has_value(); iter->next(), kvp = iter->peek()) {

auto [key, version, _value] = kvp.value();

key.remove_prefix(key_prefix.size());

int suffix = std::stoi(std::string(key));

ASSERT_EQ(suffix, tc.expected_results[count].first) << " count=" << count;

ASSERT_EQ(version, tc.expected_results[count].second)

<< " count=" << count << ", key=" << hex(key)

<< ", version=" << version.version()

<< ", expected=" << tc.expected_results[count].second.version();

count += 1;

}

ASSERT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

ASSERT_EQ(count, tc.expected_results.size());

case_index += 1;

std::cout << "Peek test case " << case_index << " passed." << std::endl;

}

}

TEST(VersionedMessageTest, GetRange2) {

auto txn_kv = std::make_shared<MemTxnKv>();

std::string_view key_prefix = "get_range2_key_";

{

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

// Insert 1000 versioned documents with version 100

for (int i = 0; i < 1000; i++) {

std::string key = fmt::format("{}{:03}", key_prefix, i);

Versionstamp version(100, 0);

VersionPB value;

value.set_version(100);

versioned_put(txn.get(), key, version, value.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

// Insert some versioned documents, with version 200

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

for (int i = 0; i < 1000; i++) {

std::string key = fmt::format("{}{:03}", key_prefix, i);

Versionstamp version(200, 0);

VersionPB value;

value.set_version(200);

versioned_put(txn.get(), key, version, value.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

}

// Iterate [0, 999] with version 200

{

std::string begin_key = fmt::format("{}{:03}", key_prefix, 0);

std::string end_key = fmt::format("{}{:03}", key_prefix, 999);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionedRangeGetOptions opts;

opts.snapshot_version = Versionstamp::max();

opts.begin_key_selector = RangeKeySelector::FIRST_GREATER_OR_EQUAL;

opts.end_key_selector = RangeKeySelector::FIRST_GREATER_THAN;

opts.batch_limit = 11;

auto iter = versioned_get_range(txn.get(), begin_key, end_key, opts);

size_t count = 0;

for (auto&& kvp = iter->next(); kvp.has_value(); kvp = iter->next()) {

auto [key, version, _value] = kvp.value();

key.remove_prefix(key_prefix.size());

int suffix = std::stoi(std::string(key));

ASSERT_EQ(suffix, 999 - count) << " count=" << count << ", key=" << key;

ASSERT_EQ(version.version(), 200)

<< " count=" << count << ", key=" << key << ", version=" << version.version();

count += 1;

}

ASSERT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

ASSERT_EQ(count, 1000) << dump_range(txn_kv.get());

}

// Iterate [0, 999] with 100

{

std::string begin_key = fmt::format("{}{:03}", key_prefix, 0);

std::string end_key = fmt::format("{}{:03}", key_prefix, 999);

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionedRangeGetOptions opts;

opts.snapshot_version = Versionstamp(150, 0);

opts.begin_key_selector = RangeKeySelector::FIRST_GREATER_OR_EQUAL;

opts.end_key_selector = RangeKeySelector::FIRST_GREATER_THAN;

opts.batch_limit = 11; // Limit to 100 results per batch

auto iter = versioned_get_range(txn.get(), begin_key, end_key, opts);

size_t count = 0;

for (auto&& kvp = iter->next(); kvp.has_value(); kvp = iter->next()) {

auto [key, version, _value] = kvp.value();

key.remove_prefix(key_prefix.size());

int suffix = std::stoi(std::string(key));

ASSERT_EQ(suffix, 999 - count) << " count=" << count << ", key=" << key;

ASSERT_EQ(version.version(), 100)

<< " count=" << count << ", key=" << key << ", version=" << version.version();

count += 1;

}

ASSERT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

ASSERT_EQ(count, 1000) << dump_range(txn_kv.get());

}

}

TEST(VersionedMessageTest, GetRangeWithRangeKeySelector) {

auto txn_kv = std::make_shared<MemTxnKv>();

std::string_view key_prefix = "get_range_key_selector_";

// Insert 9 versioned documents with version 100

for (int i = 0; i < 10; i++) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

std::string key = fmt::format("{}{:03}", key_prefix, i);

Versionstamp version(100, 0);

VersionPB value;

value.set_version(100);

versioned_put(txn.get(), key, version, value.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

struct TestCase {

std::string begin_key;

std::string end_key;

RangeKeySelector begin_selector;

RangeKeySelector end_selector;

size_t expected_count;

};

std::vector<TestCase> test_cases = {

// Test case 1: Range [0, 9] with FIRST_GREATER_OR_EQUAL and FIRST_GREATER_THAN

{fmt::format("{}{:03}", key_prefix, 0), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::FIRST_GREATER_OR_EQUAL, RangeKeySelector::FIRST_GREATER_THAN, 10},

// Test case 2: Range [0, 9) with FIRST_GREATER_OR_EQUAL and FIRST_GREATER_OR_EQUAL

{fmt::format("{}{:03}", key_prefix, 0), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::FIRST_GREATER_OR_EQUAL, RangeKeySelector::FIRST_GREATER_OR_EQUAL, 9},

// Test case 3: Range (0, 9] with FIRST_GREATER_THAN and FIRST_GREATER_THAN

{fmt::format("{}{:03}", key_prefix, 0), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::FIRST_GREATER_THAN, RangeKeySelector::FIRST_GREATER_THAN, 9},

// Test case 4: Range (0, 9) with FIRST_GREATER_THAN and FIRST_GREATER_OR_EQUAL

{fmt::format("{}{:03}", key_prefix, 0), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::FIRST_GREATER_THAN, RangeKeySelector::FIRST_GREATER_OR_EQUAL, 8},

// Test case 5: Range [0, 8) with LAST_LESS_OR_EQUAL and LAST_LESS_THAN

{fmt::format("{}{:03}", key_prefix, 0), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::LAST_LESS_OR_EQUAL, RangeKeySelector::LAST_LESS_THAN, 8},

// Test case 6: Range [0, 9) with LAST_LESS_OR_EQUAL and LAST_LESS_OR_EQUAL

{fmt::format("{}{:03}", key_prefix, 0), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::LAST_LESS_OR_EQUAL, RangeKeySelector::LAST_LESS_OR_EQUAL, 9},

// Test case 7: Range (0, 8) with LAST_LESS_THAN and LAST_LESS_THAN

{fmt::format("{}{:03}", key_prefix, 1), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::LAST_LESS_THAN, RangeKeySelector::LAST_LESS_THAN, 8},

// Test case 8: Range (0, 9) with LAST_LESS_THAN and LAST_LESS_OR_EQUAL

{fmt::format("{}{:03}", key_prefix, 1), fmt::format("{}{:03}", key_prefix, 9),

RangeKeySelector::LAST_LESS_THAN, RangeKeySelector::LAST_LESS_OR_EQUAL, 9},

};

size_t case_index = 0;

for (const auto& tc : test_cases) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionedRangeGetOptions opts;

opts.snapshot_version = Versionstamp::max();

opts.begin_key_selector = tc.begin_selector;

opts.end_key_selector = tc.end_selector;

opts.batch_limit = 11; // Limit to 100 results per batch

auto iter = versioned_get_range(txn.get(), tc.begin_key, tc.end_key, opts);

size_t count = 0;

std::string keys;

for (auto&& kvp = iter->next(); kvp.has_value(); kvp = iter->next()) {

auto [key, version, _value] = kvp.value();

keys += fmt::format("{} -> {}\n", key, version.version());

count += 1;

}

ASSERT_TRUE(iter->is_valid()); // The iterator should still be valid after the next call.

ASSERT_EQ(count, tc.expected_count)

<< keys << ", case_index=" << case_index << ", begin_key=" << tc.begin_key

<< ", end_key=" << tc.end_key

<< ", begin_selector=" << static_cast<int>(tc.begin_selector)

<< ", end_selector=" << static_cast<int>(tc.end_selector);

case_index += 1;

}

}

TEST(VersionedValueTest, RemoveAll) {

auto txn_kv = std::make_shared<MemTxnKv>();

std::string key_prefix = "remove_all_key_";

// Insert 10 versioned values

for (int i = 0; i < 10; ++i) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionPB version_pb;

version_pb.set_version(i * 10);

Versionstamp versionstamp(i * 10, 0);

versioned_put(txn.get(), key_prefix, versionstamp, version_pb.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_EQ(count_range(txn_kv.get()), 10) << dump_range(txn_kv.get());

// Remove all versioned values

{

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

versioned_remove_all(txn.get(), key_prefix);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_EQ(count_range(txn_kv.get()), 0) << dump_range(txn_kv.get());

// Insert 5 versioned values and verify

for (int i = 0; i < 5; ++i) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionPB version_pb;

version_pb.set_version(i * 100);

Versionstamp versionstamp(i * 100, 0);

versioned_put(txn.get(), key_prefix, versionstamp, version_pb.SerializeAsString());

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_EQ(count_range(txn_kv.get()), 5) << dump_range(txn_kv.get());

// Remove all versioned values again

{

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

versioned_remove_all(txn.get(), key_prefix);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

ASSERT_EQ(count_range(txn_kv.get()), 0) << dump_range(txn_kv.get());

}

TEST(VersionedValueTest, BatchGet) {

std::shared_ptr<TxnKv> txn_kv = std::make_shared<MemTxnKv>();

std::string key_prefix = "batch_get_key_";

std::vector<std::string> keys;

for (int i = 0; i < 5; ++i) {

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

VersionPB version_pb;

version_pb.set_version(i * 10);

Versionstamp versionstamp(i * 10, 0);

std::string key = fmt::format("{}{:02}", key_prefix, i);

versioned_put(txn.get(), key, versionstamp, version_pb.SerializeAsString());

keys.push_back(key);

ASSERT_EQ(txn->commit(), TxnErrorCode::TXN_OK);

}

struct TestCase {

Versionstamp snapshot_version;

std::vector<std::string> keys;

std::vector<std::optional<Versionstamp>> expected_results;

};

std::vector<TestCase> test_cases = {

// Test case 1: Get all keys with max snapshot version

{Versionstamp::max(),

keys,

{

{Versionstamp(0, 0)},

{Versionstamp(10, 0)},

{Versionstamp(20, 0)},

{Versionstamp(30, 0)},

{Versionstamp(40, 0)},

}},

// Test case 2: Get keys with a specific snapshot version (30, 0)

{Versionstamp(30, 0),

keys,

{

{Versionstamp(0, 0)},

{Versionstamp(10, 0)},

{Versionstamp(20, 0)},

{std::nullopt},

{std::nullopt},

}},

// Test case 3: Get a non-existing keys

{Versionstamp::max(),

{"non_existing_key_1", "non_existing_key_2"},

{

std::nullopt,

std::nullopt,

}},

};

for (const auto& tc : test_cases) {

std::cout << "Running test case with snapshot version: " << tc.snapshot_version.to_string()

<< std::endl;

std::vector<std::optional<std::pair<std::string, Versionstamp>>> results;

std::unique_ptr<Transaction> txn;

ASSERT_EQ(txn_kv->create_txn(&txn), TxnErrorCode::TXN_OK);

auto err = versioned_batch_get(txn.get(), tc.keys, tc.snapshot_version, &results);

ASSERT_EQ(err, TxnErrorCode::TXN_OK)

<< "Failed to batch get keys with snapshot " << tc.snapshot_version.to_string();

ASSERT_EQ(results.size(), tc.expected_results.size());

for (size_t i = 0; i < results.size(); ++i) {

if (tc.expected_results[i].has_value()) {

ASSERT_TRUE(results[i].has_value());

ASSERT_EQ(results[i]->second.version(), tc.expected_results[i]->version())

<< "Mismatch for key: " << tc.keys[i];

} else {

ASSERT_FALSE(results[i].has_value())

<< "Expected no value for key: " << tc.keys[i] << ", value: "

<< (results[i].has_value() ? results[i]->second.to_string() : "null");

}

}

}

}