vikingowl
200cdbc4bd
- Introduce `MockProvider` with configurable models, health status, generation handlers, and error simulation. - Add common test utilities and integration tests covering provider registration, model aggregation, request routing, error handling, and health refresh.
118 lines
3.5 KiB
Rust
118 lines
3.5 KiB
Rust
mod common;
|
|
|
|
use std::sync::Arc;
|
|
|
|
use futures::StreamExt;
|
|
|
|
use common::mock_provider::MockProvider;
|
|
use owlen_core::config::Config;
|
|
use owlen_core::provider::{
|
|
GenerateChunk, GenerateRequest, ModelInfo, ProviderManager, ProviderType,
|
|
};
|
|
|
|
#[allow(dead_code)]
|
|
fn base_config() -> Config {
|
|
Config {
|
|
providers: Default::default(),
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
fn make_model(name: &str, provider: &str) -> ModelInfo {
|
|
ModelInfo {
|
|
name: name.into(),
|
|
size_bytes: None,
|
|
capabilities: vec!["chat".into()],
|
|
description: Some("mock".into()),
|
|
provider: owlen_core::provider::ProviderMetadata::new(
|
|
provider,
|
|
provider,
|
|
ProviderType::Local,
|
|
false,
|
|
),
|
|
metadata: Default::default(),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn registers_providers_and_lists_ids() {
|
|
let manager = ProviderManager::default();
|
|
let provider: Arc<dyn owlen_core::provider::ModelProvider> = MockProvider::new("mock-a").into();
|
|
|
|
manager.register_provider(provider).await;
|
|
let ids = manager.provider_ids().await;
|
|
|
|
assert_eq!(ids, vec!["mock-a".to_string()]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn aggregates_models_across_providers() {
|
|
let manager = ProviderManager::default();
|
|
let provider_a = MockProvider::new("mock-a").with_models(vec![make_model("alpha", "mock-a")]);
|
|
let provider_b = MockProvider::new("mock-b").with_models(vec![make_model("beta", "mock-b")]);
|
|
|
|
manager.register_provider(provider_a.into()).await;
|
|
manager.register_provider(provider_b.into()).await;
|
|
|
|
let models = manager.list_all_models().await.unwrap();
|
|
assert_eq!(models.len(), 2);
|
|
assert!(models.iter().any(|m| m.model.name == "alpha"));
|
|
assert!(models.iter().any(|m| m.model.name == "beta"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn routes_generation_to_specific_provider() {
|
|
let manager = ProviderManager::default();
|
|
let provider = MockProvider::new("mock-gen").with_generate_handler(|_req| {
|
|
vec![
|
|
GenerateChunk::from_text("hello"),
|
|
GenerateChunk::final_chunk(),
|
|
]
|
|
});
|
|
|
|
manager.register_provider(provider.into()).await;
|
|
|
|
let request = GenerateRequest::new("mock-gen::primary");
|
|
let mut stream = manager.generate("mock-gen", request).await.unwrap();
|
|
let mut collected = Vec::new();
|
|
while let Some(chunk) = stream.next().await {
|
|
collected.push(chunk.unwrap());
|
|
}
|
|
|
|
assert_eq!(collected.len(), 2);
|
|
assert_eq!(collected[0].text.as_deref(), Some("hello"));
|
|
assert!(collected[1].is_final);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn marks_provider_unavailable_on_error() {
|
|
let manager = ProviderManager::default();
|
|
let provider = MockProvider::new("flaky")
|
|
.with_generate_error(|| owlen_core::Error::Network("boom".into()));
|
|
|
|
manager.register_provider(provider.into()).await;
|
|
let request = GenerateRequest::new("flaky::model");
|
|
let result = manager.generate("flaky", request).await;
|
|
assert!(result.is_err());
|
|
|
|
let status = manager.provider_status("flaky").await.unwrap();
|
|
assert!(matches!(
|
|
status,
|
|
owlen_core::provider::ProviderStatus::Unavailable
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn health_refresh_updates_status_cache() {
|
|
let manager = ProviderManager::default();
|
|
let provider =
|
|
MockProvider::new("healthy").with_status(owlen_core::provider::ProviderStatus::Available);
|
|
|
|
manager.register_provider(provider.into()).await;
|
|
let statuses = manager.refresh_health().await;
|
|
assert_eq!(
|
|
statuses.get("healthy"),
|
|
Some(&owlen_core::provider::ProviderStatus::Available)
|
|
);
|
|
}
|