use anyhow::Result;
use mc_mcp::infrastructure::vector_db::{DocumentToUpsert, VectorDb};
use mc_mcp::qdrant_client::qdrant::vectors_config::Config;
use mc_mcp::qdrant_client::qdrant::{
CreateCollection, Distance, VectorParams, VectorsConfig,
};
use mc_mcp::qdrant_client::Qdrant;
use mc_mcp::VectorRepository;
use std::time::Duration;
use testcontainers::core::{ContainerAsync, ContainerPort, WaitFor};
use testcontainers::runners::AsyncRunner;
use testcontainers::GenericImage;
use tokio;
use uuid::Uuid;
async fn setup_qdrant() -> Result<(VectorDb, ContainerAsync<GenericImage>)> {
let image = GenericImage::new("qdrant/qdrant", "latest")
.with_exposed_port(ContainerPort::Tcp(6334))
.with_wait_for(WaitFor::message_on_stdout("Qdrant gRPC listening on 6334"));
let container = image.start().await?;
let grpc_port = container.get_host_port_ipv4(6334).await?;
let qdrant_url = format!("http://localhost:{}", grpc_port);
let client = Qdrant::from_url(&qdrant_url).build()?;
let collection_name = format!("test_coll_{}", Uuid::new_v4().as_simple());
let vector_size: u64 = 3;
let create_collection_request = CreateCollection {
collection_name: collection_name.clone(),
vectors_config: Some(VectorsConfig {
config: Some(Config::Params(VectorParams {
size: vector_size,
distance: Distance::Cosine.into(),
hnsw_config: None,
quantization_config: None,
on_disk: None,
datatype: None,
multivector_config: None,
})),
}),
..Default::default()
};
let create_collection_result = client
.create_collection(create_collection_request)
.await;
if let Err(e) = create_collection_result {
let error_string = e.to_string();
if !error_string.contains("already exists") && !error_string.contains("already created") {
return Err(e.into());
}
log::warn!(
"Collection {} already exists, likely due to race condition. Continuing.",
collection_name
);
}
let vector_db = VectorDb::new(Box::new(client), collection_name, vector_size)?;
Ok((vector_db, container))
}
#[tokio::test]
async fn test_vector_db_new_and_initialize() -> Result<()> {
let (_vector_db, _container): (VectorDb, ContainerAsync<GenericImage>) = setup_qdrant().await?;
Ok(())
}
#[tokio::test]
async fn test_vector_db_upsert_and_search() -> Result<()> {
let (vector_db, _container): (VectorDb, ContainerAsync<GenericImage>) = setup_qdrant().await?;
let docs_to_upsert = vec![
DocumentToUpsert {
file_path: "file1.md".to_string(),
vector: vec![0.1, 0.2, 0.7],
source: Some("source_A".to_string()),
content_chunk: "This is chunk 1 from source A.".to_string(),
metadata: Some(serde_json::json!({ "section": "intro" })),
},
DocumentToUpsert {
file_path: "file2.md".to_string(),
vector: vec![0.8, 0.1, 0.1],
source: Some("source_B".to_string()),
content_chunk: "This is chunk 1 from source B.".to_string(),
metadata: None,
},
DocumentToUpsert {
file_path: "file1.md".to_string(),
vector: vec![0.2, 0.3, 0.5],
source: Some("source_A".to_string()),
content_chunk: "This is chunk 2 from source A.".to_string(),
metadata: Some(serde_json::json!({ "section": "details" })),
},
];
vector_db.upsert_documents(&docs_to_upsert).await?;
tokio::time::sleep(Duration::from_millis(500)).await;
let query_vector = vec![0.15, 0.25, 0.6];
let results = vector_db.search(query_vector.clone(), 5, Some(0.5)).await?;
log::info!("Search Results: {:?}", results);
assert!(!results.is_empty(), "Search should return results");
assert_eq!(results.len(), 2, "Expected 2 results above threshold 0.5");
let top_result = &results[0];
assert!(top_result.score > 0.5, "Score should be above threshold");
assert_eq!(top_result.file_path, "file1.md");
assert_eq!(top_result.source, Some("source_A".to_string()));
assert_eq!(top_result.content_chunk, "This is chunk 1 from source A.");
assert_eq!(
top_result.metadata,
Some(serde_json::json!({ "section": "intro" }))
);
let second_result = &results[1];
assert!(second_result.score > 0.5);
assert_eq!(second_result.file_path, "file1.md");
assert_eq!(second_result.source, Some("source_A".to_string()));
assert_eq!(
second_result.content_chunk,
"This is chunk 2 from source A."
);
assert_eq!(
second_result.metadata,
Some(serde_json::json!({ "section": "details" }))
);
let query_vector_b = vec![0.7, 0.15, 0.15];
let search_result_b = vector_db
.search(query_vector_b.clone(), 5, Some(0.5))
.await?;
assert_eq!(search_result_b.len(), 2, "Expected 2 results for query B above 0.5");
assert_eq!(search_result_b[0].file_path, "file2.md");
assert_eq!(search_result_b[0].source, Some("source_B".to_string()));
assert_eq!(
search_result_b[0].content_chunk,
"This is chunk 1 from source B."
);
assert!(search_result_b[0].metadata.is_none());
Ok(())
}
#[tokio::test]
async fn test_vector_db_new_invalid_params() -> Result<()> {
let client1 = Qdrant::from_url("http://dummy-url1").build()?;
assert!(VectorDb::new(Box::new(client1), "".to_string(), 3).is_err());
let client2 = Qdrant::from_url("http://dummy-url2").build()?;
assert!(VectorDb::new(Box::new(client2), "test".to_string(), 0).is_err());
Ok(())
}
#[tokio::test]
async fn test_vector_db_search_wrong_dimension() -> Result<()> {
let (vector_db, _container): (VectorDb, ContainerAsync<GenericImage>) = setup_qdrant().await?;
let query_vector = vec![0.1, 0.2];
let search_result = vector_db.search(query_vector, 5, None).await;
assert!(search_result.is_err());
let error_string = search_result.unwrap_err().to_string();
assert!(
error_string.contains("Query vector dimension (2) does not match collection dimension (3)"),
"Error message should indicate wrong dimension: {}", error_string
);
Ok(())
}
#[tokio::test]
async fn test_vector_db_upsert_empty() -> Result<()> {
let (vector_db, _container): (VectorDb, ContainerAsync<GenericImage>) = setup_qdrant().await?;
let empty_docs: Vec<DocumentToUpsert> = vec![];
vector_db.upsert_documents(&empty_docs).await?;
Ok(())
}