138 lines
4.2 KiB
Rust
138 lines
4.2 KiB
Rust
use brokecrew::client::{ClientError, ModelClient, OpenRouterClient};
|
|
use brokecrew::registry::ModelRegistry;
|
|
use tokio::sync::mpsc;
|
|
use wiremock::matchers::{method, path};
|
|
use wiremock::{Mock, MockServer, ResponseTemplate};
|
|
|
|
#[tokio::test]
|
|
async fn test_wiremock_sse_streaming() {
|
|
let mock_server = MockServer::start().await;
|
|
|
|
let sse_body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"from \"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"WireMock!\"}}]}\n\ndata: [DONE]\n\n";
|
|
|
|
Mock::given(method("POST"))
|
|
.and(path("/chat/completions"))
|
|
.respond_with(
|
|
ResponseTemplate::new(200)
|
|
.insert_header("content-type", "text/event-stream")
|
|
.set_body_string(sse_body),
|
|
)
|
|
.mount(&mock_server)
|
|
.await;
|
|
|
|
let client = OpenRouterClient::with_base_url("test-key".into(), mock_server.uri());
|
|
|
|
let (tx, mut rx) = mpsc::channel::<String>(10);
|
|
let collect_task = tokio::spawn(async move {
|
|
let mut chunks = Vec::new();
|
|
while let Some(c) = rx.recv().await {
|
|
chunks.push(c);
|
|
}
|
|
chunks
|
|
});
|
|
|
|
let full_reply = client
|
|
.complete_stream("mock-model", "system", "user", tx)
|
|
.await
|
|
.expect("streaming should succeed");
|
|
|
|
let chunks = collect_task.await.unwrap();
|
|
|
|
assert_eq!(full_reply, "Hello from WireMock!");
|
|
assert_eq!(chunks, vec!["Hello ", "from ", "WireMock!"]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_wiremock_401_unauthorized() {
|
|
let mock_server = MockServer::start().await;
|
|
|
|
Mock::given(method("POST"))
|
|
.and(path("/chat/completions"))
|
|
.respond_with(ResponseTemplate::new(401).set_body_string("Invalid API key"))
|
|
.mount(&mock_server)
|
|
.await;
|
|
|
|
let client = OpenRouterClient::with_base_url("bad-key".into(), mock_server.uri());
|
|
let err = client
|
|
.complete("mock-model", "sys", "user")
|
|
.await
|
|
.unwrap_err();
|
|
|
|
match err {
|
|
ClientError::Unauthorized => (),
|
|
_ => panic!("Expected ClientError::Unauthorized, got {:?}", err),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_wiremock_402_payment_required() {
|
|
let mock_server = MockServer::start().await;
|
|
|
|
Mock::given(method("POST"))
|
|
.and(path("/chat/completions"))
|
|
.respond_with(ResponseTemplate::new(402).set_body_string("Account balance depleted"))
|
|
.mount(&mock_server)
|
|
.await;
|
|
|
|
let client = OpenRouterClient::with_base_url("key".into(), mock_server.uri());
|
|
let err = client
|
|
.complete("mock-model", "sys", "user")
|
|
.await
|
|
.unwrap_err();
|
|
|
|
match err {
|
|
ClientError::PaymentRequired => (),
|
|
_ => panic!("Expected ClientError::PaymentRequired, got {:?}", err),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_wiremock_models_catalog_fetch() {
|
|
let mock_server = MockServer::start().await;
|
|
|
|
let catalog_json = r#"{
|
|
"data": [
|
|
{
|
|
"id": "qwen/qwen-2.5-coder-32b-instruct:free",
|
|
"name": "Qwen 2.5 Coder",
|
|
"context_length": 32768,
|
|
"pricing": { "prompt": "0", "completion": "0" }
|
|
},
|
|
{
|
|
"id": "meta-llama/llama-3.3-70b-instruct:free",
|
|
"name": "Llama 3.3 70B",
|
|
"context_length": 131072,
|
|
"pricing": { "prompt": "0", "completion": "0" }
|
|
}
|
|
]
|
|
}"#;
|
|
|
|
Mock::given(method("GET"))
|
|
.and(path("/models"))
|
|
.respond_with(
|
|
ResponseTemplate::new(200)
|
|
.insert_header("content-type", "application/json")
|
|
.set_body_string(catalog_json),
|
|
)
|
|
.mount(&mock_server)
|
|
.await;
|
|
|
|
let registry = ModelRegistry::new();
|
|
let count = registry
|
|
.fetch_live(&mock_server.uri())
|
|
.await
|
|
.expect("catalog fetch failed");
|
|
|
|
assert_eq!(count, 2);
|
|
assert!(
|
|
registry
|
|
.is_model_available("qwen/qwen-2.5-coder-32b-instruct:free")
|
|
.await
|
|
);
|
|
assert!(
|
|
registry
|
|
.is_model_available("meta-llama/llama-3.3-70b-instruct:free")
|
|
.await
|
|
);
|
|
assert!(!registry.is_model_available("non-existent-model").await);
|
|
}
|