From c78394214a40329bf22e447afd24806e05b7477e Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 11:12:09 -0500 Subject: [PATCH 01/10] fix: mongodb integration test --- rig-mongodb/tests/integration_tests.rs | 53 +++++++++++++++++--------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index 9b340ae2..d560c11c 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -30,6 +30,39 @@ const DATABASE_NAME: &str = "rig"; const USERNAME: &str = "riguser"; const PASSWORD: &str = "rigpassword"; +async fn create_search_index(collection: &Collection) { + let mut attempts = 0; + let max_attempts = 5; + + while attempts < max_attempts { + match collection + .create_search_index( + SearchIndexModel::builder() + .name(Some(VECTOR_SEARCH_INDEX_NAME.to_string())) + .index_type(Some(mongodb::SearchIndexType::VectorSearch)) + .definition(doc! { + "fields": [{ + "numDimensions": 1536, + "path": "embedding", + "similarity": "cosine", + "type": "vector" + }] + }) + .build(), + ) + .await + { + Ok(_) => return, + Err(_) => { + println!("Waiting for MongoDB to be ready... Attempts remaining: {}", max_attempts - attempts - 1); + sleep(Duration::from_secs(5)).await; + attempts += 1; + } + } + } + panic!("Failed to create search index after {} attempts", max_attempts); +} + #[tokio::test] async fn vector_search_test() { // Initialize OpenAI client @@ -114,24 +147,8 @@ async fn bootstrap_collection(host: String, port: u16) -> Collection Date: Tue, 17 Dec 2024 11:29:40 -0500 Subject: [PATCH 02/10] Reduce delay --- rig-mongodb/tests/integration_tests.rs | 68 +++++++++++++------------- 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index d560c11c..32403380 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -30,39 +30,6 @@ const DATABASE_NAME: &str = "rig"; const USERNAME: &str = "riguser"; const PASSWORD: &str = "rigpassword"; -async fn create_search_index(collection: &Collection) { - let mut attempts = 0; - let max_attempts = 5; - - while attempts < max_attempts { - match collection - .create_search_index( - SearchIndexModel::builder() - .name(Some(VECTOR_SEARCH_INDEX_NAME.to_string())) - .index_type(Some(mongodb::SearchIndexType::VectorSearch)) - .definition(doc! { - "fields": [{ - "numDimensions": 1536, - "path": "embedding", - "similarity": "cosine", - "type": "vector" - }] - }) - .build(), - ) - .await - { - Ok(_) => return, - Err(_) => { - println!("Waiting for MongoDB to be ready... Attempts remaining: {}", max_attempts - attempts - 1); - sleep(Duration::from_secs(5)).await; - attempts += 1; - } - } - } - panic!("Failed to create search index after {} attempts", max_attempts); -} - #[tokio::test] async fn vector_search_test() { // Initialize OpenAI client @@ -104,7 +71,7 @@ async fn vector_search_test() { .await .unwrap(); - sleep(Duration::from_secs(15)).await; + sleep(Duration::from_secs(5)).await; // Query the index let results = index @@ -124,6 +91,39 @@ async fn vector_search_test() { ) } +async fn create_search_index(collection: &Collection) { + let mut attempts = 0; + let max_attempts = 5; + + while attempts < max_attempts { + match collection + .create_search_index( + SearchIndexModel::builder() + .name(Some(VECTOR_SEARCH_INDEX_NAME.to_string())) + .index_type(Some(mongodb::SearchIndexType::VectorSearch)) + .definition(doc! { + "fields": [{ + "numDimensions": 1536, + "path": "embedding", + "similarity": "cosine", + "type": "vector" + }] + }) + .build(), + ) + .await + { + Ok(_) => return, + Err(_) => { + println!("Waiting for MongoDB to be ready... Attempts remaining: {}", max_attempts - attempts - 1); + sleep(Duration::from_secs(5)).await; + attempts += 1; + } + } + } + panic!("Failed to create search index after {} attempts", max_attempts); +} + async fn bootstrap_collection(host: String, port: u16) -> Collection { // Initialize MongoDB client let options = ClientOptions::parse(format!( From d01c8ac2874a420a16a70cfa96553825b324b8ff Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 12:28:33 -0500 Subject: [PATCH 03/10] Run linter --- rig-mongodb/tests/integration_tests.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index 32403380..3cc08655 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -94,7 +94,7 @@ async fn vector_search_test() { async fn create_search_index(collection: &Collection) { let mut attempts = 0; let max_attempts = 5; - + while attempts < max_attempts { match collection .create_search_index( @@ -115,13 +115,19 @@ async fn create_search_index(collection: &Collection) { { Ok(_) => return, Err(_) => { - println!("Waiting for MongoDB to be ready... Attempts remaining: {}", max_attempts - attempts - 1); + println!( + "Waiting for MongoDB to be ready... Attempts remaining: {}", + max_attempts - attempts - 1 + ); sleep(Duration::from_secs(5)).await; attempts += 1; } } } - panic!("Failed to create search index after {} attempts", max_attempts); + panic!( + "Failed to create search index after {} attempts", + max_attempts + ); } async fn bootstrap_collection(host: String, port: u16) -> Collection { From 6c5ed969ba709e604a7539d2cf99fc791b95991d Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 13:46:48 -0500 Subject: [PATCH 04/10] Verify index was created --- rig-mongodb/tests/integration_tests.rs | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index 3cc08655..ab01a51c 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -14,6 +14,7 @@ use testcontainers::{ GenericImage, ImageExt, }; use tokio::time::{sleep, Duration}; +use futures::StreamExt; #[derive(Embed, Clone, serde::Deserialize, serde::Serialize, Debug, PartialEq)] struct Word { @@ -124,6 +125,29 @@ async fn create_search_index(collection: &Collection) { } } } + + // Verify the index was created + attempts = 0; + while attempts < max_attempts { + let mut cursor = collection.list_search_indexes().await.unwrap(); + + while let Some(index_result) = cursor.next().await { + match index_result { + Ok(index) => { + if index.get_str("name").unwrap_or("") == VECTOR_SEARCH_INDEX_NAME { + return; + } + } + Err(e) => { + eprintln!("Error processing index: {}", e); + continue; + } + } + } + sleep(Duration::from_secs(2)).await; + attempts += 1; + } + panic!( "Failed to create search index after {} attempts", max_attempts From 0981d496f7b9b9fba48336977da620175bc2ffb3 Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 13:48:05 -0500 Subject: [PATCH 05/10] Add comment --- rig-mongodb/tests/integration_tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index 232a6859..d4f2a406 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -97,6 +97,7 @@ async fn create_search_index(collection: &Collection) { let mut attempts = 0; let max_attempts = 5; + // Create the search index while attempts < max_attempts { match collection .create_search_index( From f2875ce70330790f77515351fe391fa7dd28e7d4 Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 13:58:08 -0500 Subject: [PATCH 06/10] Fix logic --- rig-mongodb/tests/integration_tests.rs | 50 +++++++++++++------------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index d4f2a406..d8e49ff8 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -1,3 +1,4 @@ +use futures::StreamExt; use mongodb::{ bson::{self, doc}, options::ClientOptions, @@ -14,7 +15,6 @@ use testcontainers::{ GenericImage, ImageExt, }; use tokio::time::{sleep, Duration}; -use futures::StreamExt; #[derive(Embed, Clone, serde::Deserialize, serde::Serialize, Debug, PartialEq)] struct Word { @@ -116,7 +116,31 @@ async fn create_search_index(collection: &Collection) { ) .await { - Ok(_) => return, + Ok(_) => { + // Verify the index was created + let mut search_attempts = 0; + while search_attempts < max_attempts { + let mut cursor = collection.list_search_indexes().await.unwrap(); + + while let Some(index_result) = cursor.next().await { + match index_result { + Ok(index) => { + if index.get_str("name").unwrap_or("") == VECTOR_SEARCH_INDEX_NAME { + return; + } + } + Err(e) => { + println!("Error processing index: {}", e); + continue; + } + } + } + sleep(Duration::from_secs(2)).await; + search_attempts += 1; + } + + panic!("Failed to create search index",); + } Err(_) => { println!( "Waiting for MongoDB to be ready... Attempts remaining: {}", @@ -128,28 +152,6 @@ async fn create_search_index(collection: &Collection) { } } - // Verify the index was created - attempts = 0; - while attempts < max_attempts { - let mut cursor = collection.list_search_indexes().await.unwrap(); - - while let Some(index_result) = cursor.next().await { - match index_result { - Ok(index) => { - if index.get_str("name").unwrap_or("") == VECTOR_SEARCH_INDEX_NAME { - return; - } - } - Err(e) => { - eprintln!("Error processing index: {}", e); - continue; - } - } - } - sleep(Duration::from_secs(2)).await; - attempts += 1; - } - panic!( "Failed to create search index after {} attempts", max_attempts From 89f9581cf7ad440118972bab010991a89b826c5b Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 14:01:05 -0500 Subject: [PATCH 07/10] Reduce wait --- rig-mongodb/tests/integration_tests.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index d8e49ff8..a19a574e 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -61,7 +61,7 @@ async fn vector_search_test() { collection.insert_many(embeddings).await.unwrap(); // Wait for the new documents to be indexed - sleep(Duration::from_secs(30)).await; + sleep(Duration::from_secs(5)).await; // Create a vector index on our vector store. // Note: a vector index called "vector_index" must exist on the MongoDB collection you are querying. From a483b8aebc3e182e5e4974bfa470e69584ab7e8a Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 14:29:47 -0500 Subject: [PATCH 08/10] improve code --- rig-mongodb/tests/integration_tests.rs | 51 ++++++++++---------------- 1 file changed, 19 insertions(+), 32 deletions(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index a19a574e..63c00ad9 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -94,11 +94,9 @@ async fn vector_search_test() { } async fn create_search_index(collection: &Collection) { - let mut attempts = 0; let max_attempts = 5; - // Create the search index - while attempts < max_attempts { + for attempt in 0..max_attempts { match collection .create_search_index( SearchIndexModel::builder() @@ -117,45 +115,34 @@ async fn create_search_index(collection: &Collection) { .await { Ok(_) => { - // Verify the index was created - let mut search_attempts = 0; - while search_attempts < max_attempts { - let mut cursor = collection.list_search_indexes().await.unwrap(); - - while let Some(index_result) = cursor.next().await { - match index_result { - Ok(index) => { - if index.get_str("name").unwrap_or("") == VECTOR_SEARCH_INDEX_NAME { - return; - } - } - Err(e) => { - println!("Error processing index: {}", e); - continue; - } - } + // Wait for index to be available + for _ in 0..max_attempts { + let indexes = collection.list_search_indexes() + .await + .unwrap() + .collect::>() + .await; + + if indexes.iter().any(|idx| { + idx.as_ref() + .ok() + .and_then(|i| i.get_str("name").ok()) + .map_or(false, |name| name == VECTOR_SEARCH_INDEX_NAME) + }) { + return; } sleep(Duration::from_secs(2)).await; - search_attempts += 1; } - - panic!("Failed to create search index",); + panic!("Index creation verified but index not found"); } Err(_) => { - println!( - "Waiting for MongoDB to be ready... Attempts remaining: {}", - max_attempts - attempts - 1 - ); + println!("Waiting for MongoDB... {} attempts remaining", max_attempts - attempt - 1); sleep(Duration::from_secs(5)).await; - attempts += 1; } } } - panic!( - "Failed to create search index after {} attempts", - max_attempts - ); + panic!("Failed to create search index after {} attempts", max_attempts); } async fn bootstrap_collection(host: String, port: u16) -> Collection { From 83b09f18e843c8b360d3ec6fdb8ad132082c5d68 Mon Sep 17 00:00:00 2001 From: cygaar Date: Tue, 17 Dec 2024 14:41:52 -0500 Subject: [PATCH 09/10] lint --- rig-mongodb/tests/integration_tests.rs | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index 63c00ad9..64cf3b72 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -117,12 +117,13 @@ async fn create_search_index(collection: &Collection) { Ok(_) => { // Wait for index to be available for _ in 0..max_attempts { - let indexes = collection.list_search_indexes() + let indexes = collection + .list_search_indexes() .await .unwrap() .collect::>() .await; - + if indexes.iter().any(|idx| { idx.as_ref() .ok() @@ -136,13 +137,19 @@ async fn create_search_index(collection: &Collection) { panic!("Index creation verified but index not found"); } Err(_) => { - println!("Waiting for MongoDB... {} attempts remaining", max_attempts - attempt - 1); + println!( + "Waiting for MongoDB... {} attempts remaining", + max_attempts - attempt - 1 + ); sleep(Duration::from_secs(5)).await; } } } - panic!("Failed to create search index after {} attempts", max_attempts); + panic!( + "Failed to create search index after {} attempts", + max_attempts + ); } async fn bootstrap_collection(host: String, port: u16) -> Collection { From 481098c65a96a0bf23ed2c33fc460a020fba44ba Mon Sep 17 00:00:00 2001 From: cygaar Date: Wed, 18 Dec 2024 10:45:42 -0500 Subject: [PATCH 10/10] Pr feedback --- rig-mongodb/tests/integration_tests.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/rig-mongodb/tests/integration_tests.rs b/rig-mongodb/tests/integration_tests.rs index 64cf3b72..b5ef0c93 100644 --- a/rig-mongodb/tests/integration_tests.rs +++ b/rig-mongodb/tests/integration_tests.rs @@ -119,6 +119,7 @@ async fn create_search_index(collection: &Collection) { for _ in 0..max_attempts { let indexes = collection .list_search_indexes() + .name(VECTOR_SEARCH_INDEX_NAME) .await .unwrap() .collect::>() @@ -127,8 +128,19 @@ async fn create_search_index(collection: &Collection) { if indexes.iter().any(|idx| { idx.as_ref() .ok() - .and_then(|i| i.get_str("name").ok()) - .map_or(false, |name| name == VECTOR_SEARCH_INDEX_NAME) + .and_then(|i| { + // Check both name and status + let name_matches = i + .get_str("name") + .ok() + .map_or(false, |name| name == VECTOR_SEARCH_INDEX_NAME); + let status_ready = i + .get_str("status") + .ok() + .map_or(false, |status| status == "READY"); + Some(name_matches && status_ready) + }) + .unwrap_or(false) }) { return; }