Skip to content

Commit cf61c79

Browse files
authored
[CHORE][chroma] Use a collection per test rather than database per test. (#5717)
## Description of changes The client testing strategy was to create a new database each time. This PR changes that to be a collection at a time. ## Test plan CI ## Migration plan N/A ## Observability plan Tests do not flake. ## Documentation Changes N/A
1 parent b36f50a commit cf61c79

File tree

4 files changed

+151
-196
lines changed

4 files changed

+151
-196
lines changed

Cargo.lock

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

rust/chroma/src/client/chroma_http_client.rs

Lines changed: 40 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -880,7 +880,7 @@ impl ChromaHttpClient {
880880
mod tests {
881881
use super::*;
882882
use crate::client::ChromaRetryOptions;
883-
use crate::tests::with_client;
883+
use crate::tests::{unique_collection_name, with_client};
884884
use chroma_types::{EmbeddingFunctionConfiguration, EmbeddingFunctionNewConfiguration};
885885
use httpmock::{HttpMockResponse, MockServer};
886886
use std::sync::atomic::AtomicBool;
@@ -1017,11 +1017,10 @@ mod tests {
10171017
#[tokio::test]
10181018
#[test_log::test]
10191019
async fn test_live_cloud_parses_error() {
1020-
with_client(|client| async move {
1021-
client.create_collection("foo", None, None).await.unwrap();
1022-
1020+
with_client(|mut client| async move {
1021+
let collection = client.new_collection("foo").await;
10231022
let err = client
1024-
.create_collection("foo", None, None)
1023+
.create_collection(collection.name(), None, None)
10251024
.await
10261025
.unwrap_err();
10271026

@@ -1039,67 +1038,59 @@ mod tests {
10391038
#[tokio::test]
10401039
#[test_log::test]
10411040
async fn test_live_cloud_list_collections() {
1042-
with_client(|client| async move {
1043-
let collections = client.list_collections(100, None).await.unwrap();
1044-
assert!(collections.is_empty());
1045-
1046-
client.create_collection("first", None, None).await.unwrap();
1047-
1048-
client
1049-
.create_collection("second", None, None)
1050-
.await
1051-
.unwrap();
1052-
1053-
let collections = client.list_collections(100, None).await.unwrap();
1054-
assert_eq!(collections.len(), 2);
1055-
1056-
let collections = client.list_collections(1, Some(1)).await.unwrap();
1057-
assert_eq!(collections.len(), 1);
1058-
assert_eq!(collections[0].collection.name, "second");
1041+
with_client(|mut client| async move {
1042+
let first = client.new_collection("first").await;
1043+
let second = client.new_collection("second").await;
1044+
let first = first.name();
1045+
let second = second.name();
1046+
1047+
let collections = client.list_collections(1000, None).await.unwrap();
1048+
let names: std::collections::HashSet<_> = collections
1049+
.iter()
1050+
.map(|collection| collection.name().to_string())
1051+
.collect();
1052+
1053+
assert!(names.contains(first));
1054+
assert!(names.contains(second));
1055+
let positions = collections
1056+
.iter()
1057+
.enumerate()
1058+
.filter(|(_, collection)| collection.name() == first || collection.name() == second)
1059+
.collect::<Vec<_>>();
1060+
assert_eq!(positions.len(), 2);
10591061
})
10601062
.await;
10611063
}
10621064

10631065
#[tokio::test]
10641066
#[test_log::test]
10651067
async fn test_live_cloud_create_collection() {
1066-
with_client(|client| async move {
1068+
with_client(|mut client| async move {
10671069
let schema = Schema::default_with_embedding_function(
10681070
EmbeddingFunctionConfiguration::Known(EmbeddingFunctionNewConfiguration {
10691071
name: "bar".to_string(),
10701072
config: serde_json::json!({}),
10711073
}),
10721074
);
1073-
let collection = client
1074-
.create_collection("foo", Some(schema.clone()), None)
1075+
let collection1 = client.new_collection("foo").await;
1076+
let collection2 = client
1077+
.get_or_create_collection(collection1.name(), Some(schema), None)
10751078
.await
10761079
.unwrap();
1077-
assert_eq!(collection.collection.name, "foo");
1078-
1079-
let collection = client
1080-
.get_or_create_collection("foo", None, None)
1081-
.await
1082-
.unwrap();
1083-
assert_eq!(collection.schema().clone().unwrap(), schema);
1080+
assert_eq!(collection1.name(), collection2.name());
1081+
assert_eq!(collection1.schema(), collection2.schema());
10841082
})
10851083
.await;
10861084
}
10871085

10881086
#[tokio::test]
10891087
#[test_log::test]
10901088
async fn test_live_cloud_get_collection() {
1091-
with_client(|client| async move {
1092-
client
1093-
.create_collection("my_collection".to_string(), None, None)
1094-
.await
1095-
.unwrap();
1096-
1097-
let collection = client
1098-
.get_collection("my_collection".to_string())
1099-
.await
1100-
.unwrap();
1101-
1102-
assert_eq!(collection.collection.name, "my_collection");
1089+
with_client(|mut client| async move {
1090+
let collection = client.new_collection("my_collection").await;
1091+
let name = collection.name().to_string();
1092+
let collection = client.get_collection(collection.name()).await.unwrap();
1093+
assert_eq!(collection.collection.name, name);
11031094
})
11041095
.await;
11051096
}
@@ -1108,20 +1099,16 @@ mod tests {
11081099
#[test_log::test]
11091100
async fn test_live_cloud_delete_collection() {
11101101
with_client(|client| async move {
1111-
client
1112-
.create_collection("to_be_deleted".to_string(), None, None)
1113-
.await
1114-
.unwrap();
1102+
let name = unique_collection_name("to_be_deleted");
11151103

11161104
client
1117-
.delete_collection("to_be_deleted".to_string())
1105+
.create_collection(name.clone(), None, None)
11181106
.await
11191107
.unwrap();
11201108

1121-
let err = client
1122-
.get_collection("to_be_deleted".to_string())
1123-
.await
1124-
.unwrap_err();
1109+
client.delete_collection(name.clone()).await.unwrap();
1110+
1111+
let err = client.get_collection(name.clone()).await.unwrap_err();
11251112

11261113
match err {
11271114
ChromaHttpClientError::ApiError(msg, status) => {

0 commit comments

Comments
 (0)