use crate::{ReaderCachedShard, ReaderNodeHeartbeatRequest, ReaderNodeStatus};

pub fn create_test_heartbeat_request(
    addr: String,
    tiflash_addr: String,
    cached_shard_ids: Vec<u64>,
) -> ReaderNodeHeartbeatRequest {
    ReaderNodeHeartbeatRequest {
        addr,
        status: Some(ReaderNodeStatus {
            cpu_usage: 30.0,
            memory_usage: 1024,
            requests_in_cycle: 10,
            cache_hit_rate: 90.0,
        }),
        tiflash_server_addr: tiflash_addr,
        cached_shard_ids,
        cached_shards: vec![],
    }
}

pub fn create_test_heartbeat_request_with_cached_shards(
    addr: String,
    tiflash_addr: String,
    cached_shards: Vec<ReaderCachedShard>,
) -> ReaderNodeHeartbeatRequest {
    ReaderNodeHeartbeatRequest {
        addr,
        status: Some(ReaderNodeStatus {
            cpu_usage: 30.0,
            memory_usage: 1024,
            requests_in_cycle: 10,
            cache_hit_rate: 90.0,
        }),
        tiflash_server_addr: tiflash_addr,
        cached_shard_ids: cached_shards.iter().map(|shard| shard.shard_id).collect(),
        cached_shards,
    }
}
