|
| 1 | +// Copyright 2025-2026 Andrey Vasilevsky <anvanster@gmail.com> |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +//! End-to-end regression test for issue #16. |
| 5 | +//! |
| 6 | +//! The unit tests in `docs.rs` prove that chunk ids differ between sources. |
| 7 | +//! This one proves the thing the user actually reported: indexing a second |
| 8 | +//! markdown file used to delete most of the first file's chunks from RocksDB, |
| 9 | +//! so `list_doc_sources` and `search_docs` stopped returning them - while |
| 10 | +//! indexing still reported success. |
| 11 | +//! |
| 12 | +//! It drives the real `DocStore`: real RocksDB keys, real embeddings, real |
| 13 | +//! HNSW search, and a reopen to confirm what survived on disk. |
| 14 | +//! |
| 15 | +//! Needs a local model2vec model directory, which is also what the |
| 16 | +//! `--embedding-model static` server path uses. Point `CODEGRAPH_STATIC_MODEL` |
| 17 | +//! at one, or have the default `~/.codegraph/static_models/jina-code-static-256` |
| 18 | +//! in place. The test skips (with a message) when no model is available rather |
| 19 | +//! than failing, since the model is not vendored in the repo. |
| 20 | +
|
| 21 | +use codegraph_memory::{DocStore, VectorEngine}; |
| 22 | +use std::path::PathBuf; |
| 23 | +use std::sync::Arc; |
| 24 | + |
| 25 | +fn static_model_dir() -> Option<PathBuf> { |
| 26 | + let dir = match std::env::var("CODEGRAPH_STATIC_MODEL") { |
| 27 | + Ok(v) => PathBuf::from(v), |
| 28 | + Err(_) => dirs_home()? |
| 29 | + .join(".codegraph") |
| 30 | + .join("static_models") |
| 31 | + .join("jina-code-static-256"), |
| 32 | + }; |
| 33 | + dir.join("model.safetensors").exists().then_some(dir) |
| 34 | +} |
| 35 | + |
| 36 | +fn dirs_home() -> Option<PathBuf> { |
| 37 | + std::env::var_os("HOME").map(PathBuf::from) |
| 38 | +} |
| 39 | + |
| 40 | +/// Ten sections, so the document is comfortably larger than the second one. |
| 41 | +fn architecture_md() -> String { |
| 42 | + let mut md = String::from("# Architecture Guide\n\n"); |
| 43 | + for i in 1..=10 { |
| 44 | + md.push_str(&format!( |
| 45 | + "## Subsystem {i}\n\nThe subsystem-{i} component owns marker-architecture-{i} and \ |
| 46 | + is responsible for coordinating work across the graph engine. It keeps its own \ |
| 47 | + state and reports progress to the supervisor.\n\n" |
| 48 | + )); |
| 49 | + } |
| 50 | + md |
| 51 | +} |
| 52 | + |
| 53 | +/// Three sections - smaller than the guide above, which is what made the |
| 54 | +/// original data loss look intermittent. |
| 55 | +fn onboarding_md() -> String { |
| 56 | + let mut md = String::from("# Onboarding Guide\n\n"); |
| 57 | + for i in 1..=3 { |
| 58 | + md.push_str(&format!( |
| 59 | + "## Step {i}\n\nFollow step-{i} to set up your workstation; marker-onboarding-{i} \ |
| 60 | + covers the tools you need before your first change lands.\n\n" |
| 61 | + )); |
| 62 | + } |
| 63 | + md |
| 64 | +} |
| 65 | + |
| 66 | +#[test] |
| 67 | +fn indexing_a_second_source_does_not_evict_the_first() { |
| 68 | + let Some(model_dir) = static_model_dir() else { |
| 69 | + eprintln!("skipping: no static embedding model available (set CODEGRAPH_STATIC_MODEL)"); |
| 70 | + return; |
| 71 | + }; |
| 72 | + |
| 73 | + let tmp = tempfile::tempdir().expect("temp dir"); |
| 74 | + let arch_path = tmp.path().join("architecture.md"); |
| 75 | + let onboard_path = tmp.path().join("onboarding.md"); |
| 76 | + std::fs::write(&arch_path, architecture_md()).expect("write architecture.md"); |
| 77 | + std::fs::write(&onboard_path, onboarding_md()).expect("write onboarding.md"); |
| 78 | + |
| 79 | + let engine = Arc::new(VectorEngine::with_static_model(&model_dir).expect("static engine")); |
| 80 | + let db_path = tmp.path().join("docs.db"); |
| 81 | + |
| 82 | + let arch_indexed; |
| 83 | + let onboard_indexed; |
| 84 | + { |
| 85 | + let store = DocStore::new(&db_path, Arc::clone(&engine)).expect("open store"); |
| 86 | + arch_indexed = store |
| 87 | + .index_file(&arch_path, 500) |
| 88 | + .expect("index architecture.md") |
| 89 | + .len(); |
| 90 | + onboard_indexed = store |
| 91 | + .index_file(&onboard_path, 500) |
| 92 | + .expect("index onboarding.md") |
| 93 | + .len(); |
| 94 | + |
| 95 | + assert!(arch_indexed > onboard_indexed, "sanity: sizes differ"); |
| 96 | + println!("indexed architecture.md -> {arch_indexed} chunks"); |
| 97 | + println!("indexed onboarding.md -> {onboard_indexed} chunks"); |
| 98 | + |
| 99 | + let sources = store.list_sources(); |
| 100 | + println!("list_doc_sources -> {} source(s)", sources.len()); |
| 101 | + assert_eq!(sources.len(), 2, "both sources must be listed: {sources:?}"); |
| 102 | + |
| 103 | + // The first file must still have every chunk it was indexed with. |
| 104 | + let arch_source = arch_path.to_string_lossy().to_string(); |
| 105 | + let arch_stored = store.get_chunks_by_source(&arch_source).len(); |
| 106 | + println!("chunks still stored for architecture.md -> {arch_stored}"); |
| 107 | + assert_eq!( |
| 108 | + arch_stored, arch_indexed, |
| 109 | + "indexing the second file must not drop chunks from the first" |
| 110 | + ); |
| 111 | + |
| 112 | + // And it must still be findable, which is the user-visible symptom. |
| 113 | + let hits = store.search("marker-architecture-7 subsystem", 3).expect("search"); |
| 114 | + for hit in &hits { |
| 115 | + let file = std::path::Path::new(&hit.chunk.source_file); |
| 116 | + println!( |
| 117 | + "search_docs hit -> {} § {} ({:.2})", |
| 118 | + file.file_name().unwrap_or_default().to_string_lossy(), |
| 119 | + hit.chunk.title, |
| 120 | + hit.score |
| 121 | + ); |
| 122 | + } |
| 123 | + assert!( |
| 124 | + hits.iter().any(|h| h.chunk.source_file == arch_source), |
| 125 | + "search must still reach the first document" |
| 126 | + ); |
| 127 | + } |
| 128 | + |
| 129 | + // Reopen: chunk ids are RocksDB keys, so a collision would show up as |
| 130 | + // missing rows after a restart too. |
| 131 | + let store = DocStore::new(&db_path, engine).expect("reopen store"); |
| 132 | + println!( |
| 133 | + "after reopen -> {} source(s), architecture.md {} chunks, onboarding.md {} chunks", |
| 134 | + store.list_sources().len(), |
| 135 | + store.get_chunks_by_source(&arch_path.to_string_lossy()).len(), |
| 136 | + store |
| 137 | + .get_chunks_by_source(&onboard_path.to_string_lossy()) |
| 138 | + .len(), |
| 139 | + ); |
| 140 | + assert_eq!(store.list_sources().len(), 2, "both sources survive a reopen"); |
| 141 | + assert_eq!( |
| 142 | + store |
| 143 | + .get_chunks_by_source(&arch_path.to_string_lossy()) |
| 144 | + .len(), |
| 145 | + arch_indexed, |
| 146 | + "first document survives a reopen intact" |
| 147 | + ); |
| 148 | + assert_eq!( |
| 149 | + store |
| 150 | + .get_chunks_by_source(&onboard_path.to_string_lossy()) |
| 151 | + .len(), |
| 152 | + onboard_indexed, |
| 153 | + "second document survives a reopen intact" |
| 154 | + ); |
| 155 | +} |
0 commit comments