IdeA/crates/infrastructure/tests/onnx_embedder.rs

//! Tests for the real in-process ONNX-backed embedder (LOT C1b, §14.5.3), gated by
//! the `vector-onnx` feature. They exercise [`OnnxEmbedder`] and the
//! [`embedder_from_profile`] mapping *with* the feature on.
//!
//! The whole file is compiled out unless `--features vector-onnx` is set, so the
//! default dependency-free build is unaffected.
//!
//! ## What runs by default vs. behind `#[ignore]`
//!
//! - **No-network tests** (always run with the feature on): every path that
//!   short-circuits *before* `fastembed`'s `try_new`/download — empty input, an
//!   unknown model, and cheap/infallible construction. These never touch disk or
//!   the network, so they are safe in CI.
//! - **Real-download tests** (`#[ignore]`, never run by default): the ones that
//!   actually load/download the ~118 MB e5-small model. Run them on demand with
//!   `--features vector-onnx --test onnx_embedder -- --ignored`.
#![cfg(feature = "vector-onnx")]

use std::path::PathBuf;

use domain::ports::{Embedder, EmbedderError};
use domain::profile::{EmbedderProfile, EmbedderStrategy};
use infrastructure::{embedder_from_profile, OnnxEmbedder};
use uuid::Uuid;

// ---------------------------------------------------------------------------
// A unique, self-cleaning scratch dir under the OS temp dir (the project's
// established test convention — see e.g. tests/project_store.rs — rather than a
// new `tempfile` dev-dependency). It is created lazily by callers when a real
// download is involved; the no-network tests use a never-created path on purpose.
// ---------------------------------------------------------------------------

struct TempDir(PathBuf);
impl TempDir {
    fn new() -> Self {
        let p = std::env::temp_dir().join(format!("idea-onnx-{}", Uuid::new_v4()));
        std::fs::create_dir_all(&p).unwrap();
        Self(p)
    }
    fn path(&self) -> &std::path::Path {
        &self.0
    }
}
impl Drop for TempDir {
    fn drop(&mut self) {
        let _ = std::fs::remove_dir_all(&self.0);
    }
}

/// A `localOnnx` profile with the given (optional) model string and dimension.
fn onnx_profile(model: Option<&str>, dimension: usize) -> EmbedderProfile {
    EmbedderProfile::new(
        "test-onnx",
        "Test ONNX",
        EmbedderStrategy::LocalOnnx,
        model.map(str::to_string),
        None,
        None,
        dimension,
    )
    .unwrap()
}

// ===========================================================================
// No-network tests (always run with the feature on).
// ===========================================================================

#[tokio::test]
async fn onnx_unknown_model_is_unsupported() {
    // A non-empty model string outside the whitelist must surface `Unsupported`
    // *before* any `try_new`/download — so this is safe without a network. We point
    // the cache at a path that is never created to prove no download is attempted.
    let cache = std::path::Path::new("/idea-onnx-cache-that-never-exists");
    let embedder =
        OnnxEmbedder::from_profile(&onnx_profile(Some("definitely-unknown"), 384), cache);

    let err = embedder
        .embed(&["x".to_string()])
        .await
        .expect_err("unknown model must error, not embed");
    assert!(
        matches!(err, EmbedderError::Unsupported(_)),
        "unknown model ⇒ Unsupported (no download), got {err:?}"
    );
    // The cache dir must not have been created by the failed call.
    assert!(
        !cache.exists(),
        "an unknown model must not trigger any I/O / download"
    );
}

#[tokio::test]
async fn onnx_empty_input_short_circuits() {
    // An empty batch returns `Ok(vec![])` before any model load — even with a known
    // model and a never-created cache dir, so it can never download.
    let cache = std::path::Path::new("/idea-onnx-cache-that-never-exists-2");
    let embedder = OnnxEmbedder::from_profile(&onnx_profile(None, 384), cache);

    let out = embedder
        .embed(&[])
        .await
        .expect("empty input ⇒ empty output, no load");
    assert!(out.is_empty(), "empty batch ⇒ empty result");
    assert!(
        !cache.exists(),
        "empty input must not trigger any I/O / download"
    );
}

#[tokio::test]
async fn onnx_construction_is_cheap() {
    // `from_profile` is documented cheap & infallible: no panic, no I/O, and it
    // advertises the profile's id and dimension verbatim. Try both a known model
    // and an unknown one (construction never fails for either).
    let cache = std::path::Path::new("/idea-onnx-cache-that-never-exists-3");

    let known =
        OnnxEmbedder::from_profile(&onnx_profile(Some("multilingual-e5-small"), 384), cache);
    assert_eq!(known.id(), "test-onnx");
    assert_eq!(known.dimension(), 384);

    let default_model = OnnxEmbedder::from_profile(&onnx_profile(None, 384), cache);
    assert_eq!(default_model.dimension(), 384);

    let unknown = OnnxEmbedder::from_profile(&onnx_profile(Some("definitely-unknown"), 384), cache);
    // Construction still succeeds for an unknown model (the error is deferred to embed).
    assert_eq!(unknown.id(), "test-onnx");
    assert_eq!(unknown.dimension(), 384);

    assert!(
        !cache.exists(),
        "construction must perform no I/O whatsoever"
    );
}

#[tokio::test]
async fn embedder_from_profile_localonnx_is_real_engine_not_unsupported_stub() {
    // With the feature on, `localOnnx` must map to the real OnnxEmbedder, NOT the
    // Unsupported stub. We can prove "not the stub" without a download: a *known*
    // model with an empty batch returns Ok(vec![]) (the real engine short-circuits),
    // whereas the StubEmbedder would return Unsupported even for an empty batch.
    let cache = std::path::Path::new("/idea-onnx-cache-that-never-exists-4");
    let known = onnx_profile(Some("multilingual-e5-small"), 384);
    let embedder = embedder_from_profile(&known, cache).expect("localOnnx must yield an embedder");
    assert_eq!(embedder.dimension(), 384);
    let out = embedder
        .embed(&[])
        .await
        .expect("real engine short-circuits empty input to Ok(vec![])");
    assert!(
        out.is_empty(),
        "empty batch ⇒ empty result on the real engine"
    );

    // And an *unknown* model still surfaces Unsupported (deferred resolution), so the
    // mapping is the real engine in both cases (the stub would also say Unsupported,
    // but the empty-batch check above already disproves the stub for the known case).
    let unknown = onnx_profile(Some("definitely-unknown"), 384);
    let embedder =
        embedder_from_profile(&unknown, cache).expect("localOnnx must yield an embedder");
    let err = embedder
        .embed(&["x".to_string()])
        .await
        .expect_err("unknown model ⇒ Unsupported");
    assert!(
        matches!(err, EmbedderError::Unsupported(_)),
        "unknown model via mapping ⇒ Unsupported, got {err:?}"
    );
    assert!(!cache.exists(), "no I/O for these short-circuiting paths");
}

// ===========================================================================
// Real-download tests — IGNORED by default (they fetch the ~118 MB e5-small model).
// Run on demand: `cargo test -p infrastructure --features vector-onnx \
//   --test onnx_embedder -- --ignored`.
// ===========================================================================

#[tokio::test]
#[ignore = "downloads the ~118 MB e5-small ONNX model; run explicitly with --ignored"]
async fn onnx_embeds_e5_small_real_model() {
    let cache = TempDir::new();
    let embedder = OnnxEmbedder::from_profile(&onnx_profile(None, 384), cache.path());

    let texts = vec!["query: hello".to_string(), "passage: world".to_string()];
    let out = embedder
        .embed(&texts)
        .await
        .expect("real e5-small embedding must succeed");

    assert_eq!(out.len(), 2, "one vector per input");
    for v in &out {
        assert_eq!(v.len(), 384, "e5-small produces 384-dim vectors");
        let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (norm - 1.0).abs() < 1e-2,
            "fastembed L2-normalises; norm ≈ 1, got {norm}"
        );
    }

    // Deterministic across calls (model already cached after the first call).
    let again = embedder
        .embed(&texts)
        .await
        .expect("second embedding must succeed");
    assert_eq!(out, again, "embedding must be deterministic across calls");
}

#[tokio::test]
#[ignore = "loads the e5-small model to observe the dimension-mismatch validation; run with --ignored"]
async fn onnx_dimension_mismatch_is_unavailable() {
    // The profile declares 999 dimensions but e5-small produces 384; the model loads
    // fine, then per-vector validation rejects the mismatch as `Unavailable`. This
    // requires the real model load, hence `#[ignore]`.
    let cache = TempDir::new();
    let embedder = OnnxEmbedder::from_profile(&onnx_profile(Some("e5-small"), 999), cache.path());

    let err = embedder
        .embed(&["query: hello".to_string()])
        .await
        .expect_err("a profile/model dimension mismatch must error");
    assert!(
        matches!(err, EmbedderError::Unavailable(_)),
        "dimension mismatch ⇒ Unavailable, got {err:?}"
    );
}