poimen-memory/tests/it_chat_client.rs

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");
}