photoncloud-monorepo/flaredb/crates/flaredb-server/tests/test_multi_region.rs

use flaredb_server::config::{encode_namespaced_key, ServerConfig};
use flaredb_server::store::Store;
use flaredb_storage::rocks_engine::RocksEngine;
use flaredb_types::RegionMeta;
use std::collections::HashMap;
use std::sync::Arc;

fn test_peer_addrs() -> Arc<HashMap<u64, String>> {
    let mut addrs = HashMap::new();
    addrs.insert(1, "127.0.0.1:50051".to_string());
    Arc::new(addrs)
}

#[tokio::test]
async fn kv_routing_maps_keys_by_range() {
    let dir = tempfile::TempDir::new().unwrap();
    let engine = Arc::new(RocksEngine::new(dir.path().to_str().unwrap()).unwrap());
    let cfg = Arc::new(ServerConfig::default_with_namespace());
    let ns = cfg.get_namespace("default");
    let store = Arc::new(Store::new(1, engine, cfg.clone(), test_peer_addrs()));
    store
        .bootstrap_regions(vec![
            (
                RegionMeta {
                    id: 1,
                    start_key: encode_namespaced_key(ns.id, b""),
                    end_key: encode_namespaced_key(ns.id, b"m"),
                },
                vec![1],
            ),
            (
                RegionMeta {
                    id: 2,
                    start_key: encode_namespaced_key(ns.id, b"m"),
                    end_key: Vec::new(),
                },
                vec![1],
            ),
        ])
        .await
        .unwrap();

    let k1 = encode_namespaced_key(ns.id, b"a");
    let k2 = encode_namespaced_key(ns.id, b"z");
    assert_eq!(store.route_key(&k1).await, Some(1));
    assert_eq!(store.route_key(&k2).await, Some(2));
}

#[tokio::test]
async fn rejects_overlapping_regions_and_can_refresh_from_pd_meta() {
    let dir = tempfile::TempDir::new().unwrap();
    let engine = Arc::new(RocksEngine::new(dir.path().to_str().unwrap()).unwrap());
    let store = Arc::new(Store::new(
        1,
        engine.clone(),
        Arc::new(ServerConfig::default_with_namespace()),
        test_peer_addrs(),
    ));

    // Overlap detection
    let err = store
        .bootstrap_regions(vec![
            (
                RegionMeta {
                    id: 1,
                    start_key: encode_namespaced_key(1, b""),
                    end_key: encode_namespaced_key(1, b"z"),
                },
                vec![1],
            ),
            (
                RegionMeta {
                    id: 2,
                    start_key: encode_namespaced_key(1, b"y"),
                    end_key: Vec::new(),
                },
                vec![1],
            ),
        ])
        .await
        .unwrap_err();
    assert!(err.contains("overlap"));

    // PD supplies two disjoint regions; refresh replaces layout
    store
        .bootstrap_regions(vec![(
            RegionMeta {
                id: 10,
                start_key: encode_namespaced_key(1, b""),
                end_key: Vec::new(),
            },
            vec![1],
        )])
        .await
        .unwrap();
    store
        .refresh_regions(vec![
            (
                RegionMeta {
                    id: 1,
                    start_key: encode_namespaced_key(1, b""),
                    end_key: encode_namespaced_key(1, b"m"),
                },
                vec![1],
            ),
            (
                RegionMeta {
                    id: 2,
                    start_key: encode_namespaced_key(1, b"m"),
                    end_key: Vec::new(),
                },
                vec![1],
            ),
        ])
        .await
        .unwrap();

    let k1 = encode_namespaced_key(1, b"a");
    let k2 = encode_namespaced_key(1, b"x");
    assert_eq!(store.route_key(&k1).await, Some(1));
    assert_eq!(store.route_key(&k2).await, Some(2));
}