use std::collections::BTreeMap;
use async_trait::async_trait;
use crate::discovery::{Discovery, DiscoveryError, DiscoveryResult, ServiceInstance};
#[derive(Debug, Clone, Default)]
pub struct StaticDiscovery {
services: BTreeMap<String, Vec<ServiceInstance>>,
}
impl StaticDiscovery {
pub fn new() -> Self {
Self::default()
}
pub fn with_instance(mut self, instance: ServiceInstance) -> Self {
self.services
.entry(instance.service.clone())
.or_default()
.push(instance);
self
}
}
#[async_trait]
impl Discovery for StaticDiscovery {
async fn discover(&self, service: &str) -> DiscoveryResult<Vec<ServiceInstance>> {
let instances = self
.services
.get(service)
.map(|items| {
items
.iter()
.filter(|instance| instance.healthy)
.cloned()
.collect::<Vec<_>>()
})
.unwrap_or_default();
if instances.is_empty() {
Err(DiscoveryError::NoInstances {
service: service.to_string(),
})
} else {
Ok(instances)
}
}
}
#[cfg(test)]
mod tests {
use super::{Discovery, StaticDiscovery};
use crate::discovery::{InstanceEndpoint, ServiceInstance};
#[tokio::test]
async fn static_discovery_filters_unhealthy_instances() {
let endpoint = InstanceEndpoint::new("127.0.0.1", 8080).expect("endpoint");
let discovery = StaticDiscovery::new()
.with_instance(ServiceInstance::new("api", "api-1", endpoint.clone()))
.with_instance(ServiceInstance::new("api", "api-2", endpoint).with_health(false));
let instances = discovery.discover("api").await.expect("instances");
assert_eq!(instances.len(), 1);
assert_eq!(instances[0].id, "api-1");
}
}