open-standards/eoc/eoc-rs/crates/eoc-rerank/tests/vendor_rerank_http.rs

//! Integration tests for the vendor re-rank HTTP backends.

use eoc_rerank::reranker::{Candidate, Reranker};
use eoc_rerank::{CohereReranker, VoyageReranker};
use eoc_rerank::error::RerankError;
use serde_json::json;
use wiremock::matchers::{body_partial_json, header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};

fn cohere_body() -> serde_json::Value {
    json!({
        "id": "x",
        "results": [
            {"index": 1, "relevance_score": 0.95},
            {"index": 0, "relevance_score": 0.20},
        ],
        "meta": {}
    })
}

fn voyage_body() -> serde_json::Value {
    json!({
        "object": "list",
        "data": [
            {"index": 1, "relevance_score": 0.99},
            {"index": 0, "relevance_score": 0.10},
        ],
        "model": "rerank-2",
        "usage": {"total_tokens": 1}
    })
}

#[tokio::test]
async fn cohere_happy_path_reorders_candidates() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v2/rerank"))
        .and(header("authorization", "Bearer sk-test"))
        .and(body_partial_json(json!({
            "model": "rerank-english-v3.0",
            "query": "what is EOC?",
            "documents": ["irrelevant filler", "Energy-optimized compute"]
        })))
        .respond_with(ResponseTemplate::new(200).set_body_json(cohere_body()))
        .mount(&server)
        .await;

    let r = CohereReranker::new("sk-test", "rerank-english-v3.0")
        .expect("ok")
        .with_endpoint(format!("{}/v2/rerank", server.uri()));

    let scored = r
        .rerank(
            "what is EOC?",
            &[
                Candidate::new("a", "irrelevant filler"),
                Candidate::new("b", "Energy-optimized compute"),
            ],
        )
        .await
        .expect("rerank ok");

    assert_eq!(scored.len(), 2);
    assert_eq!(scored[0].candidate.id, "b");
    assert_eq!(scored[0].rank, 1);
    assert!(scored[0].score > scored[1].score);
}

#[tokio::test]
async fn cohere_401_no_retry() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v2/rerank"))
        .respond_with(ResponseTemplate::new(401))
        .expect(1)
        .mount(&server)
        .await;
    let r = CohereReranker::new("sk-bad", "rerank-english-v3.0")
        .expect("ok")
        .with_endpoint(format!("{}/v2/rerank", server.uri()));
    let res = r.rerank("q", &[Candidate::new("a", "x")]).await;
    assert!(matches!(res, Err(RerankError::InvalidApiKey)));
}

#[tokio::test]
async fn cohere_429_then_succeeds() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v2/rerank"))
        .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "0"))
        .up_to_n_times(1)
        .mount(&server)
        .await;
    Mock::given(method("POST"))
        .and(path("/v2/rerank"))
        .respond_with(ResponseTemplate::new(200).set_body_json(cohere_body()))
        .mount(&server)
        .await;
    let r = CohereReranker::new("k", "rerank-english-v3.0")
        .expect("ok")
        .with_endpoint(format!("{}/v2/rerank", server.uri()))
        .with_max_retries(3);
    let scored = r
        .rerank(
            "q",
            &[Candidate::new("a", "x"), Candidate::new("b", "y")],
        )
        .await
        .expect("eventually ok");
    assert_eq!(scored.len(), 2);
}

#[tokio::test]
async fn voyage_happy_path() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v1/rerank"))
        .respond_with(ResponseTemplate::new(200).set_body_json(voyage_body()))
        .mount(&server)
        .await;
    let r = VoyageReranker::new("k", "rerank-2")
        .expect("ok")
        .with_endpoint(format!("{}/v1/rerank", server.uri()));
    let scored = r
        .rerank(
            "q",
            &[Candidate::new("a", "x"), Candidate::new("b", "y")],
        )
        .await
        .expect("ok");
    assert_eq!(scored[0].candidate.id, "b");
    assert_eq!(scored[0].rank, 1);
}

#[tokio::test]
async fn voyage_429_exhausts_budget() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v1/rerank"))
        .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "0"))
        .mount(&server)
        .await;
    let r = VoyageReranker::new("k", "rerank-2")
        .expect("ok")
        .with_endpoint(format!("{}/v1/rerank", server.uri()))
        .with_max_retries(1);
    let res = r.rerank("q", &[Candidate::new("a", "x")]).await;
    assert!(matches!(res, Err(RerankError::RateLimited { .. })));
}