use anthropic_async::AnthropicConfig;
use anthropic_async::Client;
use anthropic_async::types::ModelListParams;
use serde_json::json;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
#[tokio::test]
async fn test_models_list_retries_on_429_then_success() {
let server = MockServer::start().await;
let count = Arc::new(AtomicUsize::new(0));
let count_clone = Arc::clone(&count);
Mock::given(method("GET"))
.and(path("/v1/models"))
.respond_with(move |_req: &wiremock::Request| {
let i = count_clone.fetch_add(1, Ordering::SeqCst);
if i == 0 {
ResponseTemplate::new(429)
.insert_header("retry-after-ms", "100")
.set_body_json(
json!({"error": {"message": "Rate limit", "type": "rate_limit_error"}}),
)
} else {
ResponseTemplate::new(200).set_body_json(json!({"data": [], "has_more": false}))
}
})
.mount(&server)
.await;
let cfg = AnthropicConfig::new()
.with_api_base(server.uri())
.with_api_key("test");
let client = Client::with_config(cfg);
let _ = client
.models()
.list(&ModelListParams::default())
.await
.unwrap();
assert!(count.load(Ordering::SeqCst) >= 2);
}
#[tokio::test]
async fn test_models_get_retries_on_500_then_success() {
let server = MockServer::start().await;
let count = Arc::new(AtomicUsize::new(0));
let count_clone = Arc::clone(&count);
Mock::given(method("GET"))
.and(path("/v1/models/claude-foo"))
.respond_with(move |_req: &wiremock::Request| {
let i = count_clone.fetch_add(1, Ordering::SeqCst);
if i == 0 {
ResponseTemplate::new(500).set_body_string("server error")
} else {
ResponseTemplate::new(200).set_body_json(json!({
"id": "claude-foo",
"created_at": "2024-06-01T12:00:00Z",
"display_name": "Claude Foo",
"type": "model"
}))
}
})
.mount(&server)
.await;
let cfg = AnthropicConfig::new()
.with_api_base(server.uri())
.with_api_key("test");
let client = Client::with_config(cfg);
let m = client.models().get("claude-foo").await.unwrap();
assert_eq!(m.id, "claude-foo");
}