// client_local_hot_cache_test.cpp
#include "client.h"
#include "local_hot_cache.h"
#include "replica.h"
#include "utils.h"

#include <glog/logging.h>
#include <gtest/gtest.h>

#include <atomic>
#include <chrono>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

// Network headers for getting local IP (Linux/Unix)
#include <ifaddrs.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <arpa/inet.h>

namespace mooncake {
namespace testing {

// Test data size constants
constexpr size_t K1_KB = 1024;
constexpr size_t K2_KB = 2048;

// Async wait time constant
constexpr int K_ASYNC_WAIT_TIME_MS = 100;

// Helper function to get local IP address
// Priority: 1) Environment variable MC_TEST_LOCAL_IP, 2) System call to get first non-loopback IP
std::string GetLocalIpAddress() {
    // Try to get IP from environment variable first (for testing flexibility)
    const char* env_ip = std::getenv("MC_TEST_LOCAL_IP");
    if (env_ip && strlen(env_ip) > 0) {
        return std::string(env_ip);
    }
    
    // Fallback: get local IP using system calls
    // Get the first non-loopback IPv4 address
    struct ifaddrs *ifaddr;
    struct ifaddrs *ifa;
    std::string ip = "127.0.0.1";  // Default fallback
    
    if (getifaddrs(&ifaddr) == -1) {
        return ip;
    }
    
    for (ifa = ifaddr; ifa != nullptr; ifa = ifa->ifa_next) {
        if (ifa->ifa_addr == nullptr) {
            continue;
        }
        
        // Look for IPv4 address
        if (ifa->ifa_addr->sa_family == AF_INET) {
            // Skip loopback interface
            if (strcmp(ifa->ifa_name, "lo") == 0) {
                continue;
            }
            
            // Check if interface is UP
            if (!(ifa->ifa_flags & IFF_UP)) {
                continue;
            }
            
            struct sockaddr_in* sin = static_cast<struct sockaddr_in*>(
                static_cast<void*>(ifa->ifa_addr));
            char host[NI_MAXHOST];
            int result = getnameinfo(static_cast<struct sockaddr*>(
                static_cast<void*>(sin)), sizeof(struct sockaddr_in),
                                     host, NI_MAXHOST, nullptr, 0, NI_NUMERICHOST);
            if (result == 0) {
                ip = std::string(host);
                break;  // Use the first non-loopback interface
            }
        }
    }
    
    freeifaddrs(ifaddr);
    return ip;
}

class LocalHotCacheTest : public ::testing::Test {
protected:
    void SetUp() override {
        google::InitGoogleLogging("LocalHotCacheTest");
        FLAGS_logtostderr = 1;
    }

    void TearDown() override {
        google::ShutdownGoogleLogging();
    }

    // Helper to create a slice with test data
    Slice CreateSlice(size_t size, char fill_char = 'A') {
        std::vector<char> data(size, fill_char);
        Slice slice;
        slice.ptr = data.data();
        slice.size = size;
        test_data_.push_back(std::move(data));  // Keep data alive
        return slice;
    }

    // Helper to verify slice data in cache
    void VerifySliceData(HotMemBlock* block, size_t expected_size, char expected_char) {
        ASSERT_NE(block, nullptr);
        ASSERT_EQ(block->size, expected_size);
        const char* data = static_cast<const char*>(block->addr);
        for (size_t i = 0; i < expected_size; ++i) {
            EXPECT_EQ(data[i], expected_char) << "Data mismatch at offset " << i;
        }
    }

    // Helper to setup client with hot cache enabled and mount segment
    struct TestClientContext {
        std::shared_ptr<Client> client;
        void* segment_ptr;
        size_t segment_size;
        const char* original_env;
    };

    TestClientContext SetupTestClientWithHotCache() {
        TestClientContext ctx;
        ctx.original_env = std::getenv("LOCAL_HOT_CACHE_SIZE");
        setenv("LOCAL_HOT_CACHE_SIZE", "33554432", 1);  // 32MB = 2 blocks
        
        std::string local_ip = GetLocalIpAddress();
        std::string local_hostname = local_ip + ":12345";
        
        auto client_opt = Client::Create(local_hostname, "redis://localhost:6379",
                                         "tcp", std::nullopt, "localhost:50051");
        if (!client_opt.has_value()) {
            return ctx;  // Return empty context
        }
        
        ctx.client = client_opt.value();
        if (!ctx.client->IsHotCacheEnabled()) {
            return ctx;  // Return context with client but hot cache not enabled
        }
        
        ctx.segment_size = 64 * 1024 * 1024;  // 64MB
        ctx.segment_ptr = allocate_buffer_allocator_memory(ctx.segment_size);
        if (ctx.segment_ptr == nullptr) {
            return ctx;  // Return context without segment
        }
        
        auto mount_result = ctx.client->MountSegment(ctx.segment_ptr, ctx.segment_size);
        if (!mount_result.has_value()) {
            free_memory("", ctx.segment_ptr);
            ctx.segment_ptr = nullptr;
            return ctx;  // Return context with unmounted segment
        }
        
        return ctx;
    }

    void CleanupTestClient(TestClientContext& ctx) {
        if (ctx.client && ctx.segment_ptr) {
            ctx.client->UnmountSegment(ctx.segment_ptr, ctx.segment_size);
            free_memory("", ctx.segment_ptr);
        } else if (ctx.segment_ptr) {
            free_memory("", ctx.segment_ptr);
        }
        if (ctx.original_env) {
            setenv("LOCAL_HOT_CACHE_SIZE", ctx.original_env, 1);
        } else {
            unsetenv("LOCAL_HOT_CACHE_SIZE");
        }
    }

    void PutTestData(Client* client, const std::string& key, const std::string& data) {
        ReplicateConfig config;
        config.replica_num = 1;
        std::vector<char> put_buffer(data.size());
        std::memcpy(put_buffer.data(), data.data(), data.size());
        std::vector<Slice> put_slices;
        put_slices.emplace_back(Slice{put_buffer.data(), data.size()});
        auto put_result = client->Put(key, put_slices, config);
        // If Put fails, the test will fail when trying to Get the data
    }

private:
    std::vector<std::vector<char>> test_data_;  // Keep test data alive
};

// Test LocalHotCache construction
TEST_F(LocalHotCacheTest, Construction) {
    const size_t cache_size = 64 * 1024 * 1024;  // 64MB = 4 blocks
    LocalHotCache cache(cache_size);
    
    EXPECT_GT(cache.GetCacheSize(), 0);
    EXPECT_EQ(cache.GetCacheSize(), 4);  // 64MB / 16MB = 4 blocks
}

// Test LocalHotCache with zero size
TEST_F(LocalHotCacheTest, ZeroSizeCache) {
    LocalHotCache cache(0);
    EXPECT_EQ(cache.GetCacheSize(), 0);
}

// Test PutHotSlice basic functionality
TEST_F(LocalHotCacheTest, PutHotSliceBasic) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    LocalHotCache cache(cache_size);
    
    const size_t slice_size = K1_KB;
    Slice slice = CreateSlice(slice_size, 'X');
    
    EXPECT_TRUE(cache.PutHotSlice("test_key_0", slice));
    EXPECT_TRUE(cache.HasHotSlice("test_key_0"));
    
    HotMemBlock* block = cache.GetHotSlice("test_key_0");
    VerifySliceData(block, slice_size, 'X');
}

// Test PutHotSlice with multiple keys
TEST_F(LocalHotCacheTest, PutHotSliceMultipleKeys) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    LocalHotCache cache(cache_size);
    
    // Put first key
    Slice slice1 = CreateSlice(K1_KB, 'A');
    EXPECT_TRUE(cache.PutHotSlice("key1", slice1));
    
    // Put second key
    Slice slice2 = CreateSlice(K2_KB, 'B');
    EXPECT_TRUE(cache.PutHotSlice("key2", slice2));
    
    EXPECT_TRUE(cache.HasHotSlice("key1"));
    EXPECT_TRUE(cache.HasHotSlice("key2"));
    
    HotMemBlock* block1 = cache.GetHotSlice("key1");
    VerifySliceData(block1, K1_KB, 'A');
    
    HotMemBlock* block2 = cache.GetHotSlice("key2");
    VerifySliceData(block2, K2_KB, 'B');
}

// Test PutHotSlice with existing key (touch LRU)
TEST_F(LocalHotCacheTest, PutHotSliceTouchExisting) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    LocalHotCache cache(cache_size);
    
    Slice slice = CreateSlice(K1_KB, 'X');
    EXPECT_TRUE(cache.PutHotSlice("test_key", slice));
    
    // Put same key again (should just touch LRU, not copy data)
    Slice slice2 = CreateSlice(K1_KB, 'Y');
    EXPECT_TRUE(cache.PutHotSlice("test_key", slice2));
    
    // Data should still be 'X' (not updated)
    HotMemBlock* block = cache.GetHotSlice("test_key");
    VerifySliceData(block, K1_KB, 'X');
}

// Test PutHotSlice with invalid parameters
TEST_F(LocalHotCacheTest, PutHotSliceInvalidParams) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    LocalHotCache cache(cache_size);
    
    // Test with null pointer
    Slice null_slice;
    null_slice.ptr = nullptr;
    null_slice.size = K1_KB;
    EXPECT_FALSE(cache.PutHotSlice("key", null_slice));
    
    // Test with zero size
    Slice zero_slice;
    zero_slice.ptr = malloc(K1_KB);
    zero_slice.size = 0;
    EXPECT_FALSE(cache.PutHotSlice("key", zero_slice));
    free(zero_slice.ptr);
    
    // Test with size larger than block size
    Slice large_slice;
    large_slice.ptr = malloc(17 * 1024 * 1024);  // 17MB > 16MB
    large_slice.size = 17 * 1024 * 1024;
    EXPECT_FALSE(cache.PutHotSlice("key", large_slice));
    free(large_slice.ptr);
}

// Test LRU eviction behavior
TEST_F(LocalHotCacheTest, LRUEviction) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    LocalHotCache cache(cache_size);
    
    // Fill cache with 2 blocks
    Slice slice1 = CreateSlice(K1_KB, 'A');
    EXPECT_TRUE(cache.PutHotSlice("key1", slice1));
    
    Slice slice2 = CreateSlice(K1_KB, 'B');
    EXPECT_TRUE(cache.PutHotSlice("key2", slice2));
    
    // Add third key, should evict key1 (LRU)
    Slice slice3 = CreateSlice(K1_KB, 'C');
    EXPECT_TRUE(cache.PutHotSlice("key3", slice3));
    
    EXPECT_FALSE(cache.HasHotSlice("key1"));  // Should be evicted
    EXPECT_TRUE(cache.HasHotSlice("key2"));
    EXPECT_TRUE(cache.HasHotSlice("key3"));
}

// Test GetHotSlice updates LRU
TEST_F(LocalHotCacheTest, GetHotSliceUpdatesLRU) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    LocalHotCache cache(cache_size);
    
    Slice slice1 = CreateSlice(K1_KB, 'A');
    cache.PutHotSlice("key1", slice1);
    
    Slice slice2 = CreateSlice(K1_KB, 'B');
    cache.PutHotSlice("key2", slice2);
    
    // Access key1, should move it to front
    HotMemBlock* block1 = cache.GetHotSlice("key1");
    ASSERT_NE(block1, nullptr);
    
    // Add third key, should evict key2 (not key1, since key1 was accessed)
    Slice slice3 = CreateSlice(K1_KB, 'C');
    cache.PutHotSlice("key3", slice3);
    
    EXPECT_TRUE(cache.HasHotSlice("key1"));  // Should still be there
    EXPECT_FALSE(cache.HasHotSlice("key2"));  // Should be evicted
    EXPECT_TRUE(cache.HasHotSlice("key3"));
}

// Test GetHotSlice with non-existent key
TEST_F(LocalHotCacheTest, GetHotSliceMiss) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    LocalHotCache cache(cache_size);
    
    HotMemBlock* block = cache.GetHotSlice("non_existent_key");
    EXPECT_EQ(block, nullptr);
}

// Test LocalHotCacheHandler basic functionality
TEST_F(LocalHotCacheTest, LocalHotCacheHandlerBasic) {
    const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
    auto cache = std::make_shared<LocalHotCache>(cache_size);
    LocalHotCacheHandler handler(cache);
    
    const size_t slice_size = K1_KB;
    std::vector<char> data(slice_size, 'Z');
    Slice slice;
    slice.ptr = data.data();
    slice.size = slice_size;
    
    // Submit task
    EXPECT_TRUE(handler.SubmitPutTask("async_key", slice));
    
    // Wait a bit for async processing
    std::this_thread::sleep_for(std::chrono::milliseconds(K_ASYNC_WAIT_TIME_MS));
    
    // Verify data was cached
    EXPECT_TRUE(cache->HasHotSlice("async_key"));
    HotMemBlock* block = cache->GetHotSlice("async_key");
    VerifySliceData(block, slice_size, 'Z');
}

// Test LocalHotCacheHandler with null cache
TEST_F(LocalHotCacheTest, LocalHotCacheHandlerNullCache) {
    LocalHotCacheHandler handler(nullptr);
    
    std::vector<char> data(K1_KB, 'X');
    Slice slice;
    slice.ptr = data.data();
    slice.size = K1_KB;
    
    // Should return false since hot_cache_ is null
    EXPECT_FALSE(handler.SubmitPutTask("key", slice));
}

// Test concurrent access to LocalHotCache
TEST_F(LocalHotCacheTest, ConcurrentAccess) {
    const size_t cache_size = 128 * 1024 * 1024;  // 128MB = 8 blocks
    LocalHotCache cache(cache_size);
    
    const int num_threads = 4;
    const int keys_per_thread = 2;  // Total: 8 keys, cache has 8 blocks
    std::vector<std::thread> threads;
    std::atomic<int> successful_puts(0);
    std::atomic<int> successful_gets(0);
    
    // Each thread puts and gets keys
    for (int t = 0; t < num_threads; ++t) {
        threads.emplace_back([&cache, t, keys_per_thread, &successful_puts, &successful_gets]() {
            for (int i = 0; i < keys_per_thread; ++i) {
                std::string key = "thread_" + std::to_string(t) + "_key_" + std::to_string(i);
                std::vector<char> data(K1_KB, static_cast<char>('A' + t));
                Slice slice;
                slice.ptr = data.data();
                slice.size = K1_KB;
                
                // Put the key
                EXPECT_TRUE(cache.PutHotSlice(key, slice));
                successful_puts++;
                
                // Get it back - should succeed since cache has enough capacity
                HotMemBlock* block = cache.GetHotSlice(key);
                EXPECT_NE(block, nullptr) << "Block should not be nullptr with sufficient cache capacity";
                successful_gets++;
                
                // Verify data integrity
                const char* cached_data = static_cast<const char*>(block->addr);
                EXPECT_EQ(cached_data[0], static_cast<char>('A' + t));
            }
        });
    }
    
    // Wait for all threads
    for (auto& thread : threads) {
        thread.join();
    }
    
    // All puts and gets should succeed with sufficient capacity
    EXPECT_EQ(successful_puts.load(), num_threads * keys_per_thread);
    EXPECT_EQ(successful_gets.load(), num_threads * keys_per_thread);
}


/**
 * Note: The following private functions are tested indirectly through public APIs:
 * - Client::updateReplicaDescriptorFromCache: Tested via Client::Get and Client::BatchGet
 * - Client::ProcessSlicesAsync: Tested via Client::Get and Client::BatchGet
 * - Client::InitLocalHotCache: Tested via Client::Create with LOCAL_HOT_CACHE_SIZE env var
 * Note: These tests require a running master server. Run:
 * redis-server --port 6379 --protected-mode no --bind 0.0.0.0
 * ./mooncake-store/src/mooncake_master --rpc-port 50051
 */

/**
 * Test InitLocalHotCache via Client::Create and IsHotCacheEnabled
 * Note: InitLocalHotCache is a private function called by Client::Create
 */

 // Test 1: Valid environment variable - hot cache should be enabled
TEST_F(LocalHotCacheTest, InitLocalHotCacheViaClientCreate_ValidSize) {
    // Save original env var
    const char* original_env = std::getenv("LOCAL_HOT_CACHE_SIZE");
    
    setenv("LOCAL_HOT_CACHE_SIZE", "33554432", 1);  // 32MB = 2 blocks (16MB each)
    
    auto client_opt = Client::Create(
        "localhost",
        "redis://localhost:6379",  // Redis metadata server on port 6379
        "tcp",
        std::nullopt,
        "localhost:50051");
    if (client_opt.has_value()) {
        EXPECT_TRUE(client_opt.value()->IsHotCacheEnabled());
        EXPECT_EQ(client_opt.value()->GetLocalHotCacheBlockCount(), 2);
    }

    if (original_env) {
        setenv("LOCAL_HOT_CACHE_SIZE", original_env, 1);
    } else {
        unsetenv("LOCAL_HOT_CACHE_SIZE");
    }
}

// Test 2: Invalid environment variable - InitLocalHotCache returns INVALID_PARAMS
TEST_F(LocalHotCacheTest, InitLocalHotCacheViaClientCreate_InvalidEnvVar) {
    // Save original env var
    const char* original_env = std::getenv("LOCAL_HOT_CACHE_SIZE");
    
    setenv("LOCAL_HOT_CACHE_SIZE", "invalid", 1);
    
    auto client_opt = Client::Create(
        "localhost",
        "redis://localhost:6379",  // Redis metadata server on port 6379
        "tcp",
        std::nullopt,
        "localhost:50051");
    if (client_opt.has_value()) {
        EXPECT_FALSE(client_opt.value()->IsHotCacheEnabled());
        EXPECT_EQ(client_opt.value()->GetLocalHotCacheBlockCount(), 0);
    }
    
    if (original_env) {
        setenv("LOCAL_HOT_CACHE_SIZE", original_env, 1);
    } else {
        unsetenv("LOCAL_HOT_CACHE_SIZE");
    }
}

// Test 3: Zero value - InitLocalHotCache returns INVALID_PARAMS
TEST_F(LocalHotCacheTest, InitLocalHotCacheViaClientCreate_ZeroSize) {
    // Save original env var
    const char* original_env = std::getenv("LOCAL_HOT_CACHE_SIZE");
    
    setenv("LOCAL_HOT_CACHE_SIZE", "0", 1);
    
    auto client_opt = Client::Create(
        "localhost",
        "redis://localhost:6379",  // Redis metadata server on port 6379
        "tcp",
        std::nullopt,
        "localhost:50051");
    if (client_opt.has_value()) {
        EXPECT_FALSE(client_opt.value()->IsHotCacheEnabled());
        EXPECT_EQ(client_opt.value()->GetLocalHotCacheBlockCount(), 0);
    }
    
    if (original_env) {
        setenv("LOCAL_HOT_CACHE_SIZE", original_env, 1);
    } else {
        unsetenv("LOCAL_HOT_CACHE_SIZE");
    }
}

// Test 4: Negative value - InitLocalHotCache returns INVALID_PARAMS
TEST_F(LocalHotCacheTest, InitLocalHotCacheViaClientCreate_NegativeSize) {
    // Save original env var
    const char* original_env = std::getenv("LOCAL_HOT_CACHE_SIZE");
    
    setenv("LOCAL_HOT_CACHE_SIZE", "-1", 1);
    
    auto client_opt = Client::Create(
        "localhost",
        "redis://localhost:6379",  // Redis metadata server on port 6379
        "tcp",
        std::nullopt,
        "localhost:50051");
    if (client_opt.has_value()) {
        EXPECT_FALSE(client_opt.value()->IsHotCacheEnabled());
        EXPECT_EQ(client_opt.value()->GetLocalHotCacheBlockCount(), 0);
    }
    
    if (original_env) {
        setenv("LOCAL_HOT_CACHE_SIZE", original_env, 1);
    } else {
        unsetenv("LOCAL_HOT_CACHE_SIZE");
    }
}

// Test 5: Size less than 1 block (16MB) - hot cache should not be enabled
TEST_F(LocalHotCacheTest, InitLocalHotCacheViaClientCreate_LessThanOneBlock) {
    // Save original env var
    const char* original_env = std::getenv("LOCAL_HOT_CACHE_SIZE");
    
    setenv("LOCAL_HOT_CACHE_SIZE", "8388608", 1);  // 8MB < 16MB (1 block)
    
    auto client_opt = Client::Create(
        "localhost",
        "redis://localhost:6379",  // Redis metadata server on port 6379
        "tcp",
        std::nullopt,
        "localhost:50051");
    // If client creation succeeded, hot cache should be disabled
    // (because 8MB < 16MB results in 0 blocks, causing InitLocalHotCache to reset hot_cache_)
    if (client_opt.has_value()) {
        EXPECT_FALSE(client_opt.value()->IsHotCacheEnabled());
        EXPECT_EQ(client_opt.value()->GetLocalHotCacheBlockCount(), 0);
    }
    // Restore original env var
    if (original_env) {
        setenv("LOCAL_HOT_CACHE_SIZE", original_env, 1);
    } else {
        unsetenv("LOCAL_HOT_CACHE_SIZE");
    }
}

// Test 6: No environment variable - hot cache should be disabled (valid case)
TEST_F(LocalHotCacheTest, InitLocalHotCacheViaClientCreate_NoEnvVar) {
    // Save original env var
    const char* original_env = std::getenv("LOCAL_HOT_CACHE_SIZE");
    
    unsetenv("LOCAL_HOT_CACHE_SIZE");
    
    auto client_opt = Client::Create(
        "localhost",
        "redis://localhost:6379",  // Redis metadata server on port 6379
        "tcp",
        std::nullopt,
        "localhost:50051");
    if (client_opt.has_value()) {
        EXPECT_FALSE(client_opt.value()->IsHotCacheEnabled());
        EXPECT_EQ(client_opt.value()->GetLocalHotCacheBlockCount(), 0);
    }
    
    if (original_env) {
        setenv("LOCAL_HOT_CACHE_SIZE", original_env, 1);
    } else {
        unsetenv("LOCAL_HOT_CACHE_SIZE");
    }
}

/**
 * Test updateReplicaDescriptorFromCache and ProcessSlicesAsync indirectly
 * through Client::Get and Client::BatchGet
 * IMPORTANT: With a single client, data is stored locally, so:
 * - ProcessSlicesAsync won't cache local transfers (transport_endpoint_ == local_hostname_)
 * - updateReplicaDescriptorFromCache won't find cached data (nothing was cached)
 * - we need multiple clients (one Put, another Get from non-local) to test cache hit scenarios
 * These tests verify the functions work correctly when hot cache is disabled,
 * but cannot fully test cache hit scenarios with a single client setup.
 */
TEST_F(LocalHotCacheTest, GetWithHotCacheEnabledButLocalTransfer) {
    auto ctx = SetupTestClientWithHotCache();
    ASSERT_GT(ctx.client->GetLocalHotCacheBlockCount(), 0);
    
    const std::string test_key = "test_hot_cache_key";
    const std::string test_data = "Test data for hot cache";
    PutTestData(ctx.client.get(), test_key, test_data);
    
    std::vector<char> get_buffer(test_data.size());
    std::vector<Slice> get_slices;
    get_slices.emplace_back(Slice{get_buffer.data(), test_data.size()});
    
    auto get_result = ctx.client->Get(test_key, get_slices);
    ASSERT_TRUE(get_result.has_value()) << "Get should succeed";
    
    ASSERT_GE(get_slices.size(), 1) << "At least one slice should exist";
    ASSERT_GE(get_slices[0].size, test_data.size()) << "Slice size should be at least expected data size";
    EXPECT_EQ(std::memcmp(get_slices[0].ptr, test_data.data(), test_data.size()), 0)
        << "Retrieved data should match original data";
    
    CleanupTestClient(ctx);
}

TEST_F(LocalHotCacheTest, BatchGetWithHotCacheEnabledButLocalTransfer) {
    auto ctx = SetupTestClientWithHotCache();
    
    const std::vector<std::string> test_keys = {"batch_key1", "batch_key2"};
    const std::vector<std::string> test_data = {"Data1", "Data2"};
    
    for (size_t i = 0; i < test_keys.size(); ++i) {
        PutTestData(ctx.client.get(), test_keys[i], test_data[i]);
    }
    
    std::unordered_map<std::string, std::vector<Slice>> batch_slices;
    std::vector<std::vector<char>> batch_buffers;
    for (size_t i = 0; i < test_keys.size(); ++i) {
        batch_buffers.emplace_back(test_data[i].size());
        batch_slices[test_keys[i]].emplace_back(
            Slice{batch_buffers.back().data(), test_data[i].size()});
    }
    
    auto batch_get_result = ctx.client->BatchGet(test_keys, batch_slices);
    ASSERT_EQ(batch_get_result.size(), test_keys.size());
    for (const auto& result : batch_get_result) {
        ASSERT_TRUE(result.has_value()) << "BatchGet should succeed";
    }
    
    for (size_t i = 0; i < test_keys.size(); ++i) {
        auto slices_it = batch_slices.find(test_keys[i]);
        ASSERT_NE(slices_it, batch_slices.end());
        ASSERT_GE(slices_it->second.size(), 1);
        const auto& slice = slices_it->second[0];
        ASSERT_GE(slice.size, test_data[i].size());
        EXPECT_EQ(std::memcmp(slice.ptr, test_data[i].data(), test_data[i].size()), 0);
    }
    
    CleanupTestClient(ctx);
}

}  // namespace testing
}  // namespace mooncake

int main(int argc, char** argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}