use mem_llm::ChatClient; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; #[tokio::test] async fn a1_sends_apikey_header() { let mock_server = MockServer::start().await; Mock::given(method("POST")) .and(path("/chat/completions")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "choices": [{ "message": { "role": "assistant", "content": "pong" } }], "usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15 } }))) .mount(&mock_server) .await; let client = ChatClient::new(&mock_server.uri(), "test-key", "qwen2.5:3b-instruct").unwrap(); let result = client.complete("system", "user", 2048).await; assert!(result.is_ok()); // Verify the mock received exactly 1 request let reqs = mock_server.received_requests().await.unwrap(); assert_eq!(reqs.len(), 1); let req = &reqs[0]; // Assert apikey header is present assert!( req.headers.get("apikey").is_some(), "apikey header should be present" ); // Assert no Authorization header assert!( req.headers.get("Authorization").is_none(), "Authorization header should not be present" ); } #[tokio::test] async fn a2_no_tools_field() { let mock_server = MockServer::start().await; Mock::given(method("POST")) .and(path("/chat/completions")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "choices": [{ "message": { "role": "assistant", "content": "test" } }], "usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15 } }))) .mount(&mock_server) .await; let client = ChatClient::new(&mock_server.uri(), "test-key", "qwen2.5:3b-instruct").unwrap(); let _result = client.complete("system", "user", 2048).await; let reqs = mock_server.received_requests().await.unwrap(); let req = &reqs[0]; let body_str = String::from_utf8(req.body.clone()).unwrap(); let body_json: serde_json::Value = serde_json::from_str(&body_str).unwrap(); // Assert "tools" key is completely absent, not just empty assert!( !body_json.as_object().unwrap().contains_key("tools"), "tools key should not be present in request body" ); } #[tokio::test] async fn a3_retries_5xx() { let mock_server = MockServer::start().await; // First two requests return 503, third returns 200 Mock::given(method("POST")) .and(path("/chat/completions")) .respond_with( ResponseTemplate::new(503).set_body_json(serde_json::json!({ "error": "Service Unavailable" })), ) .up_to_n_times(2) .mount(&mock_server) .await; Mock::given(method("POST")) .and(path("/chat/completions")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "choices": [{ "message": { "role": "assistant", "content": "success" } }], "usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15 } }))) .mount(&mock_server) .await; let client = ChatClient::new(&mock_server.uri(), "test-key", "qwen2.5:3b-instruct").unwrap(); let result = client.complete("system", "user", 2048).await; assert!(result.is_ok()); let completion = result.unwrap(); assert_eq!(completion.text, "success"); // Verify we got exactly 3 requests (2 failures + 1 success) let reqs = mock_server.received_requests().await.unwrap(); assert_eq!( reqs.len(), 3, "Should have made 3 requests (2 retries + success)" ); } #[tokio::test] async fn a4_does_not_retry_4xx() { let mock_server = MockServer::start().await; Mock::given(method("POST")) .and(path("/chat/completions")) .respond_with(ResponseTemplate::new(400).set_body_json(serde_json::json!({ "error": { "message": "[] is too short - 'messages'" } }))) .mount(&mock_server) .await; let client = ChatClient::new(&mock_server.uri(), "test-key", "qwen2.5:3b-instruct").unwrap(); let result = client.complete("system", "user", 2048).await; assert!(result.is_err()); let error_msg = format!("{:?}", result.err().unwrap()); assert!( error_msg.contains("Client error") || error_msg.contains("400"), "Error should mention client error or 400 status" ); // Verify we made exactly 1 request (no retries) let reqs = mock_server.received_requests().await.unwrap(); assert_eq!( reqs.len(), 1, "Should have made exactly 1 request (no retries for 4xx)" ); } #[tokio::test] async fn a5_timeout_is_configurable() { let mock_server = MockServer::start().await; // Set up a mock that delays for 5 seconds Mock::given(method("POST")) .and(path("/chat/completions")) .respond_with( ResponseTemplate::new(200) .set_delay(std::time::Duration::from_secs(5)) .set_body_json(serde_json::json!({ "choices": [{ "message": { "role": "assistant", "content": "slow response" } }], "usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15 } })), ) .mount(&mock_server) .await; // Create client with 500ms timeout let client = ChatClient::new(&mock_server.uri(), "test-key", "qwen2.5:3b-instruct") .unwrap() .with_timeout(std::time::Duration::from_millis(500)); let result = client.complete("system", "user", 2048).await; assert!(result.is_err()); let error_msg = format!("{:?}", result); assert!( error_msg.to_lowercase().contains("timeout") || error_msg.to_lowercase().contains("request failed"), "Error should indicate a timeout, got: {}", error_msg ); } #[tokio::test] #[ignore] async fn a6_live_smoke() { // This test requires live gateway connectivity // Run with: cargo test --test it_chat_client -- --ignored let api_key = std::env::var("MEM_API_KEY").expect("MEM_API_KEY env var required"); let client = ChatClient::new("https://api.riotpiao.com/v1", api_key, "qwen2.5:3b-instruct").unwrap(); let result = client .complete("You are a helpful assistant.", "Reply with exactly: pong", 100) .await; assert!(result.is_ok(), "Live gateway should respond"); let completion = result.unwrap(); assert!( completion.text.to_lowercase().contains("pong"), "Response should contain 'pong': {}", completion.text ); assert!(completion.usage.total_tokens > 0, "Should report token usage"); }