AkurAI-Notes / crates/notes-api/tests/embed_test.rs

Latest change 705e1ce035ea4b9fb4c413dd38f550650185b571 - release: v0.0.5 by Claude Code

use akurai_json::Value;
use akurai_storage::BTree;
use notes_api::auth::idp_client::IdpConfig;
use notes_api::embed::{embed_query_with_key, maybe_reembed, EmbedStore};
use notes_api::NotesState;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::path::PathBuf;
use std::sync::{Mutex, OnceLock};
use std::thread;

fn setup_db(name: &str) -> (BTree, String) {
    let db_path = format!("_test_embed_{}.db", name);
    let _ = std::fs::remove_file(&db_path);
    let tree = BTree::open(&db_path).unwrap();
    (tree, db_path)
}

fn env_lock() -> std::sync::MutexGuard<'static, ()> {
    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
    LOCK.get_or_init(|| Mutex::new(())).lock().unwrap()
}

struct EnvGuard(Vec<(&'static str, Option<String>)>);

impl EnvGuard {
    fn capture(keys: &[&'static str]) -> Self {
        Self(
            keys.iter()
                .map(|key| (*key, std::env::var(key).ok()))
                .collect(),
        )
    }
}

impl Drop for EnvGuard {
    fn drop(&mut self) {
        for (key, value) in &self.0 {
            match value {
                Some(value) => std::env::set_var(key, value),
                None => std::env::remove_var(key),
            }
        }
    }
}

fn test_idp_config() -> IdpConfig {
    IdpConfig {
        issuer: "https://auth.olibuijr.com".into(),
        internal_url: "http://127.0.0.1:3500".into(),
        client_id: "notes-test".into(),
        client_secret: "secret".into(),
        redirect_uri: "https://akurai-notes.olibuijr.com/auth/callback".into(),
        allow: vec!["user@example.com".into()],
    }
}

fn spawn_embed_server(expected_auth: Option<&str>) -> (String, thread::JoinHandle<String>) {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = format!("http://{}", listener.local_addr().unwrap());
    let expected_auth = expected_auth.map(str::to_string);
    let handle = thread::spawn(move || {
        let (mut stream, _) = listener.accept().unwrap();
        let mut request = Vec::new();
        let mut buf = [0_u8; 1024];
        loop {
            let read = stream.read(&mut buf).unwrap();
            if read == 0 {
                break;
            }
            request.extend_from_slice(&buf[..read]);
            let text = String::from_utf8_lossy(&request);
            let Some(head_end) = text.find("\r\n\r\n") else {
                continue;
            };
            let content_length = text
                .lines()
                .find_map(|line| {
                    let (name, value) = line.split_once(':')?;
                    if name.eq_ignore_ascii_case("Content-Length") {
                        value.trim().parse::<usize>().ok()
                    } else {
                        None
                    }
                })
                .unwrap_or(0);
            if request.len() >= head_end + 4 + content_length {
                break;
            }
        }

        let request_text = String::from_utf8_lossy(&request).to_string();
        match expected_auth {
            Some(token) => {
                let expected = format!("\r\nAuthorization: Bearer {token}\r\n");
                assert!(
                    request_text.contains(&expected),
                    "missing bearer header in request: {request_text}"
                );
            }
            None => assert!(
                !request_text.contains("\r\nAuthorization: "),
                "unexpected authorization header in request: {request_text}"
            ),
        }

        let body = r#"{"data":[{"embedding":[0.5,0.25,-0.75]}]}"#;
        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            body.len(),
            body
        );
        stream.write_all(response.as_bytes()).unwrap();
        request_text
    });
    (addr, handle)
}

fn make_note_key(id: u64) -> Vec<u8> {
    let mut k = b"notes:".to_vec();
    k.extend_from_slice(&id.to_be_bytes());
    k
}

fn insert_note(db: &mut BTree, id: u64, title: &str, body: &str, owner: &str) {
    let now = NotesState::now();
    let record = Value::Object(vec![
        ("id".into(), Value::Int(id as i64)),
        ("title".into(), Value::Str(title.into())),
        ("body".into(), Value::Str(body.into())),
        ("ownerSub".into(), Value::Str(owner.into())),
        ("createdAt".into(), Value::Int(now)),
        ("updatedAt".into(), Value::Int(now)),
    ]);
    db.insert(&make_note_key(id), record.to_json().as_bytes())
        .unwrap();
    db.commit().unwrap();
}

#[test]
fn embed_store_round_trips_vector() {
    let (mut db, db_path) = setup_db("roundtrip");
    let mut store = EmbedStore { db: &mut db };

    let vec: Vec<f32> = vec![0.1, 0.2, -0.3, 0.5];
    store.store(1, &vec).unwrap();

    let loaded = store.load(1);
    assert!(loaded.is_some());
    let loaded = loaded.unwrap();
    assert_eq!(loaded.len(), 4);
    for (a, b) in vec.iter().zip(loaded.iter()) {
        assert!((a - b).abs() < 1e-6);
    }

    store.remove(1).unwrap();
    assert!(store.load(1).is_none());

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn embed_store_returns_none_for_missing_key() {
    let (mut db, db_path) = setup_db("missing");
    let mut store = EmbedStore { db: &mut db };
    assert!(store.load(999).is_none());
    std::fs::remove_file(&db_path).ok();
}

#[test]
fn rank_notes_returns_only_owned_notes_in_order() {
    let (mut db, db_path) = setup_db("rank");

    insert_note(&mut db, 1, "Meeting", "Budget review meeting", "alice");
    insert_note(&mut db, 2, "Shopping", "Buy groceries", "alice");
    insert_note(&mut db, 3, "Other", "Bob's note", "bob");

    let mut store = EmbedStore { db: &mut db };
    let v1 = vec![0.9, 0.1, 0.0];
    store.store(1, &v1).unwrap();
    let v2 = vec![0.1, 0.9, 0.0];
    store.store(2, &v2).unwrap();
    let v3 = vec![0.85, 0.15, 0.0];
    store.store(3, &v3).unwrap();

    let qv = vec![0.8, 0.2, 0.0];

    let ranked = store.rank_notes(&qv, &["alice"], 5);
    assert_eq!(
        ranked.len(),
        2,
        "should return 2 alice notes, got {ranked:?}"
    );
    assert_eq!(ranked[0].0, 1, "note 1 should rank first (closest to qv)");
    assert_eq!(ranked[1].0, 2, "note 2 should rank second");

    let ranked = store.rank_notes(&qv, &["alice"], 1);
    assert_eq!(ranked.len(), 1);

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn maybe_reembed_skips_when_url_is_none() {
    let (mut db, db_path) = setup_db("skip");
    insert_note(&mut db, 1, "Test", "Hello world", "user");

    maybe_reembed(&mut db, &None, &None, &None, 1, "Hello world");

    let mut store = EmbedStore { db: &mut db };
    assert!(store.load(1).is_none());

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn embed_query_with_key_sends_bearer_token() {
    let (embed_url, handle) = spawn_embed_server(Some("router-secret"));

    let vec = embed_query_with_key(
        "secure query",
        &embed_url,
        "test-embed-model",
        Some("router-secret"),
    )
    .unwrap();

    assert_eq!(vec, vec![0.5, 0.25, -0.75]);
    let request = handle.join().unwrap();
    assert!(request.contains(r#""model":"test-embed-model""#));
    assert!(request.contains(r#""input":"secure query""#));
}

#[test]
fn embed_query_with_key_skips_bearer_for_blank_token() {
    let (embed_url, handle) = spawn_embed_server(None);

    let vec =
        embed_query_with_key("plain query", &embed_url, "test-embed-model", Some("   ")).unwrap();

    assert_eq!(vec, vec![0.5, 0.25, -0.75]);
    let _ = handle.join().unwrap();
}

#[test]
fn notes_state_reads_embed_env_values() {
    let _lock = env_lock();
    let _env = EnvGuard::capture(&[
        "AKURAI_NOTES_EMBED_URL",
        "AKURAI_NOTES_EMBED_MODEL",
        "AKURAI_NOTES_EMBED_API_KEY",
    ]);
    std::env::set_var(
        "AKURAI_NOTES_EMBED_URL",
        "http://127.0.0.1:4219/v1/embeddings",
    );
    std::env::set_var("AKURAI_NOTES_EMBED_MODEL", "intfloat/multilingual-e5-small");
    std::env::set_var("AKURAI_NOTES_EMBED_API_KEY", "router-secret");

    let db_path = format!("_test_state_embed_{}.db", std::process::id());
    let _ = std::fs::remove_file(&db_path);
    let state = NotesState::new(&db_path, PathBuf::from("site/frontend"), test_idp_config())
        .expect("state should build");

    assert_eq!(
        state.embed_url.as_deref(),
        Some("http://127.0.0.1:4219/v1/embeddings")
    );
    assert_eq!(
        state.embed_model.as_deref(),
        Some("intfloat/multilingual-e5-small")
    );
    assert_eq!(state.embed_api_key.as_deref(), Some("router-secret"));

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn notes_state_ignores_blank_embed_env_values() {
    let _lock = env_lock();
    let _env = EnvGuard::capture(&[
        "AKURAI_NOTES_EMBED_URL",
        "AKURAI_NOTES_EMBED_MODEL",
        "AKURAI_NOTES_EMBED_API_KEY",
    ]);
    std::env::set_var("AKURAI_NOTES_EMBED_URL", "   ");
    std::env::set_var("AKURAI_NOTES_EMBED_MODEL", "");
    std::env::set_var("AKURAI_NOTES_EMBED_API_KEY", "\t");

    let db_path = format!("_test_state_embed_blank_{}.db", std::process::id());
    let _ = std::fs::remove_file(&db_path);
    let state = NotesState::new(&db_path, PathBuf::from("site/frontend"), test_idp_config())
        .expect("state should build");

    assert_eq!(state.embed_url, None);
    assert_eq!(state.embed_model, None);
    assert_eq!(state.embed_api_key, None);

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn substring_search_via_prefix_scan() {
    let (mut db, db_path) = setup_db("substr");
    insert_note(
        &mut db,
        1,
        "Meeting minutes",
        "Discussed Q2 targets",
        "alice",
    );
    insert_note(&mut db, 2, "Grocery list", "Milk and eggs", "bob");
    insert_note(&mut db, 3, "Budget review", "Q1 financial results", "alice");
    insert_note(&mut db, 4, "Random", "Nothing important", "alice");

    let prefix = b"notes:";
    let mut bound = prefix.to_vec();
    bound.push(0xff);

    let mut results: Vec<(u64, String)> = Vec::new();
    if let Ok(entries) = db.range(prefix, &bound) {
        for entry in entries {
            let (key, val) = entry;
            let id_bytes: [u8; 8] = key[prefix.len()..].try_into().unwrap();
            let id: u64 = u64::from_be_bytes(id_bytes);
            let Ok(v) = akurai_json::parse(&String::from_utf8_lossy(&val)) else {
                continue;
            };
            if v.get("ownerSub").and_then(|x| x.as_str()) != Some("alice") {
                continue;
            }
            let title = v.get("title").and_then(|x| x.as_str()).unwrap_or("");
            let body = v.get("body").and_then(|x| x.as_str()).unwrap_or("");
            let combined = format!("{} {}", title, body).to_lowercase();
            if combined.contains("q") {
                let what = format!("{} — {}", title, body);
                results.push((id, what));
            }
        }
    }

    assert_eq!(results.len(), 2, "should find 2 alice notes with 'q'");
    let titles: Vec<&str> = results.iter().map(|r| r.1.as_str()).collect();
    assert!(titles[0].contains("Q2") || titles[1].contains("Q2"));
    assert!(titles[0].contains("Q1") || titles[1].contains("Q1"));

    std::fs::remove_file(&db_path).ok();
}

#[test]
#[ignore]
fn create_note_stores_embedding() {
    let (mut db, _db_path) = setup_db("create_embed");
    let note_id = 1i64;
    let body = "Test note content for embedding";

    let embed_url = std::env::var("AKURAI_NOTES_EMBED_URL").ok();
    let embed_model = std::env::var("AKURAI_NOTES_EMBED_MODEL").ok();

    maybe_reembed(&mut db, &embed_url, &embed_model, &None, note_id, body);

    let mut store = EmbedStore { db: &mut db };
    let loaded = store.load(note_id);
    assert!(
        loaded.is_some(),
        "embedding should be stored when embed server is up"
    );
    let vec = loaded.unwrap();
    assert!(!vec.is_empty(), "embedding should have non-zero dimensions");
}

#[test]
#[ignore]
fn update_note_reembeds_only_when_body_changes() {}

#[test]
#[ignore]
fn search_is_owner_scoped() {}