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

use flaredb_proto::kvrpc::kv_raw_server::KvRaw;
use flaredb_proto::kvrpc::{RawGetRequest, RawPutRequest};
use flaredb_server::config::{parse_namespace_modes, ConsistencyMode, Config, NamespaceManager};
use flaredb_server::service::KvServiceImpl;
use flaredb_server::store::Store;
use flaredb_storage::rocks_engine::RocksEngine;
use flaredb_types::RegionMeta;
use std::collections::HashMap;
use std::sync::Arc;
use tonic::Request;

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)
}

#[test]
fn default_config_is_strong() {
    let config = Config::default();
    let ns_mgr = NamespaceManager::from_config(&config);
    assert_eq!(ns_mgr.get_namespace("default").mode, ConsistencyMode::Strong);
    // unknown namespace falls back to default
    assert_eq!(ns_mgr.get_namespace("unknown").mode, ConsistencyMode::Strong);
}

#[test]
fn parses_eventual_namespace_mode() {
    let parsed = parse_namespace_modes(&vec!["tenantA=eventual".to_string()]).unwrap();
    let mut config = Config::default();
    config.namespace_modes = parsed;
    let ns_mgr = NamespaceManager::from_config(&config);
    assert_eq!(ns_mgr.get_namespace("tenantA").mode, ConsistencyMode::Eventual);
}

#[test]
fn encodes_and_decodes_ts_value() {
    let val = b"hello".to_vec();
    let ts = 123u64;
    let encoded = flaredb_server::config::encode_value_with_ts(ts, &val);
    let (decoded_ts, decoded_val) = flaredb_server::config::decode_value_with_ts(&encoded);
    assert_eq!(ts, decoded_ts);
    assert_eq!(val, decoded_val);
}

#[tokio::test]
async fn eventual_mode_overwrites_with_newer_ts() {
    let cfg_map = parse_namespace_modes(&vec!["logs=eventual".to_string()]).unwrap();
    let mut config = Config::default();
    config.namespace_modes = cfg_map;
    let namespace_manager = Arc::new(NamespaceManager::from_config(&config));
    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(config),
        namespace_manager.clone(),
        test_peer_addrs(),
    ));
    store
        .bootstrap_regions(vec![(
            RegionMeta {
                id: 1,
                start_key: Vec::new(),
                end_key: Vec::new(),
            },
            vec![1],
        )])
        .await
        .unwrap();
    let service = KvServiceImpl::new(engine.clone(), namespace_manager, store.clone());

    // First write
    service
        .raw_put(Request::new(RawPutRequest {
            namespace: "logs".into(),
            key: b"k".to_vec(),
            value: b"v1".to_vec(),
        }))
        .await
        .unwrap();
    // Second write (newer ts)
    tokio::time::sleep(std::time::Duration::from_millis(1)).await;
    service
        .raw_put(Request::new(RawPutRequest {
            namespace: "logs".into(),
            key: b"k".to_vec(),
            value: b"v2".to_vec(),
        }))
        .await
        .unwrap();

    let resp = service
        .raw_get(Request::new(RawGetRequest {
            namespace: "logs".into(),
            key: b"k".to_vec(),
        }))
        .await
        .unwrap()
        .into_inner();
    assert!(resp.found);
    assert_eq!(resp.value, b"v2");
}

#[tokio::test]
async fn lww_resolves_partition_on_replay() {
    // Simulate two replicas reconciling: existing older value vs incoming newer value.
    let existing = flaredb_server::config::encode_value_with_ts(1, b"old");
    let incoming = flaredb_server::config::encode_value_with_ts(5, b"new");
    let (existing_ts, _) = flaredb_server::config::decode_value_with_ts(&existing);
    let (incoming_ts, _) = flaredb_server::config::decode_value_with_ts(&incoming);
    assert!(incoming_ts > existing_ts);

    // Apply LWW resolution the same way on_ready does.
    let result = if incoming_ts >= existing_ts {
        incoming
    } else {
        existing
    };
    assert_eq!(flaredb_server::config::decode_value_with_ts(&result).1, b"new");
}

#[test]
fn mode_switch_preserves_id_and_applies() {
    let config = Config::default();
    let ns_mgr = NamespaceManager::from_config(&config);
    let default = ns_mgr.get_namespace("default");
    assert_eq!(default.mode, ConsistencyMode::Strong);
    let id = default.id;

    // Switch to eventual then back to strong; id should remain stable and explicit should be true.
    let eventual = ns_mgr
        .set_namespace_mode("default", ConsistencyMode::Eventual)
        .unwrap();
    assert_eq!(eventual.id, id);
    assert_eq!(eventual.mode, ConsistencyMode::Eventual);
    assert!(eventual.explicit);

    let strong = ns_mgr
        .set_namespace_mode("default", ConsistencyMode::Strong)
        .unwrap();
    assert_eq!(strong.id, id);
    assert_eq!(strong.mode, ConsistencyMode::Strong);
    assert!(strong.explicit);
}

#[test]
fn detects_mode_mismatch_between_nodes() {
    let config_a = Config::default();
    let ns_mgr_a = NamespaceManager::from_config(&config_a);
    let config_b = Config::default();
    let ns_mgr_b = NamespaceManager::from_config(&config_b);

    // Diverge the mode for a namespace
    ns_mgr_a
        .set_namespace_mode("logs", ConsistencyMode::Eventual)
        .unwrap();
    ns_mgr_b
        .set_namespace_mode("logs", ConsistencyMode::Strong)
        .unwrap();

    let diffs = ns_mgr_a.diff_modes(&ns_mgr_b);
    assert_eq!(diffs.len(), 1);
    let diff = &diffs[0];
    assert_eq!(diff.namespace, "logs");
    assert_eq!(diff.self_mode, ConsistencyMode::Eventual);
    assert_eq!(diff.other_mode, ConsistencyMode::Strong);
}

#[test]
fn reserved_namespace_auto_creation() {
    use flaredb_server::config::RESERVED_NAMESPACES;

    let config = Config::default();
    let ns_mgr = NamespaceManager::from_config(&config);

    // Reserved namespaces should be automatically created with Strong mode
    for &reserved in RESERVED_NAMESPACES {
        let ns = ns_mgr.get_namespace(reserved);
        assert_eq!(
            ns.mode,
            ConsistencyMode::Strong,
            "reserved namespace '{}' should be Strong",
            reserved
        );
        assert!(ns_mgr.is_reserved(reserved));
    }
}

#[test]
fn reserved_namespace_mode_change_blocked() {
    let config = Config::default();
    let ns_mgr = NamespaceManager::from_config(&config);

    // Attempting to change mode of reserved namespace should fail
    let result = ns_mgr.set_namespace_mode("iam", ConsistencyMode::Eventual);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("reserved namespace"),
        "error should mention reserved namespace: {}",
        err
    );

    // Setting the same mode (Strong) should succeed
    let result = ns_mgr.set_namespace_mode("iam", ConsistencyMode::Strong);
    assert!(result.is_ok());
}