1+ package tech .amikos .chromadb .v2 ;
2+
3+ import org .junit .BeforeClass ;
4+ import org .junit .Test ;
5+ import tech .amikos .chromadb .v2 .AuthProvider ;
6+ import tech .amikos .chromadb .v2 .Collection ;
7+ import tech .amikos .chromadb .v2 .ChromaClient ;
8+ import tech .amikos .chromadb .v2 .*;
9+
10+ import java .util .*;
11+ import java .util .concurrent .*;
12+ import java .util .stream .Collectors ;
13+ import java .util .stream .IntStream ;
14+
15+ import static org .junit .Assert .*;
16+
17+ public class V2StressTest {
18+ private static ChromaClient client ;
19+
20+ @ BeforeClass
21+ public static void setup () {
22+ String chromaUrl = System .getenv ("CHROMA_URL" );
23+ if (chromaUrl == null ) {
24+ chromaUrl = "http://localhost:8000" ;
25+ }
26+
27+ client = ChromaClient .builder ()
28+ .serverUrl (chromaUrl )
29+ .auth (AuthProvider .none ())
30+ .connectTimeout (60 )
31+ .readTimeout (60 )
32+ .writeTimeout (60 )
33+ .tenant ("default_tenant" )
34+ .database ("default_database" )
35+ .build ();
36+ }
37+
38+ @ Test
39+ public void testLargeScale () throws Exception {
40+ String collectionName = "stress_test_" + UUID .randomUUID ().toString ().substring (0 , 8 );
41+ Collection collection = client .createCollection (collectionName );
42+
43+ // Add 10,000 records in batches of 100
44+ for (int batch = 0 ; batch < 100 ; batch ++) {
45+ List <String > ids = new ArrayList <>();
46+ List <List <Float >> embeddings = new ArrayList <>();
47+ List <Map <String , Object >> metadatas = new ArrayList <>();
48+
49+ for (int i = 0 ; i < 100 ; i ++) {
50+ int recordId = batch * 100 + i ;
51+ ids .add ("id_" + recordId );
52+
53+ // Create random embedding
54+ List <Float > embedding = new ArrayList <>();
55+ Random rand = new Random (recordId );
56+ for (int j = 0 ; j < 384 ; j ++) {
57+ embedding .add (rand .nextFloat ());
58+ }
59+ embeddings .add (embedding );
60+
61+ metadatas .add (Map .of (
62+ "batch" , batch ,
63+ "index" , i ,
64+ "category" , "category_" + (recordId % 10 )
65+ ));
66+ }
67+
68+ collection .add ()
69+ .ids (ids )
70+ .embeddings (embeddings )
71+ .metadatas (metadatas )
72+ .execute ();
73+
74+ if (batch % 10 == 0 ) {
75+ System .out .println ("Added " + ((batch + 1 ) * 100 ) + " records" );
76+ }
77+ }
78+
79+ assertEquals (10000 , collection .count ());
80+ System .out .println ("Successfully added 10,000 records" );
81+
82+ // Test queries
83+ Random rand = new Random ();
84+ List <Float > queryEmbedding = IntStream .range (0 , 384 )
85+ .mapToObj (i -> rand .nextFloat ())
86+ .collect (Collectors .toList ());
87+
88+ QueryResponse result = collection .query ()
89+ .queryEmbeddings (Arrays .asList (queryEmbedding ))
90+ .nResults (100 )
91+ .include (Include .METADATAS , Include .DISTANCES )
92+ .execute ();
93+
94+ assertEquals (1 , result .getIds ().size ());
95+ assertEquals (100 , result .getIds ().get (0 ).size ());
96+
97+ client .deleteCollection (collectionName );
98+ }
99+
100+ @ Test
101+ public void testConcurrentOperations () throws Exception {
102+ String collectionName = "concurrent_test_" + UUID .randomUUID ().toString ().substring (0 , 8 );
103+ Collection collection = client .createCollection (collectionName );
104+
105+ ExecutorService executor = Executors .newFixedThreadPool (10 );
106+ List <Future <Boolean >> futures = new ArrayList <>();
107+
108+ // Submit 100 concurrent operations
109+ for (int i = 0 ; i < 100 ; i ++) {
110+ final int taskId = i ;
111+ futures .add (executor .submit (() -> {
112+ try {
113+ String id = "concurrent_" + taskId ;
114+ List <Float > embedding = IntStream .range (0 , 384 )
115+ .mapToObj (j -> (float ) (taskId * 0.01 ))
116+ .collect (Collectors .toList ());
117+
118+ collection .add ()
119+ .ids (Arrays .asList (id ))
120+ .embeddings (Arrays .asList (embedding ))
121+ .execute ();
122+ return true ;
123+ } catch (Exception e ) {
124+ e .printStackTrace ();
125+ return false ;
126+ }
127+ }));
128+ }
129+
130+ // Wait for all operations to complete
131+ for (Future <Boolean > future : futures ) {
132+ assertTrue (future .get (30 , TimeUnit .SECONDS ));
133+ }
134+
135+ executor .shutdown ();
136+ executor .awaitTermination (1 , TimeUnit .MINUTES );
137+
138+ assertEquals (100 , collection .count ());
139+ client .deleteCollection (collectionName );
140+ }
141+
142+ @ Test
143+ public void testMemoryEfficiency () throws Exception {
144+ String collectionName = "memory_test_" + UUID .randomUUID ().toString ().substring (0 , 8 );
145+ Collection collection = client .createCollection (collectionName );
146+
147+ Runtime runtime = Runtime .getRuntime ();
148+ long initialMemory = runtime .totalMemory () - runtime .freeMemory ();
149+
150+ // Add records in a memory-efficient way
151+ int totalRecords = 5000 ;
152+ int batchSize = 50 ;
153+
154+ for (int batch = 0 ; batch < totalRecords / batchSize ; batch ++) {
155+ List <String > ids = new ArrayList <>();
156+ List <List <Float >> embeddings = new ArrayList <>();
157+
158+ for (int i = 0 ; i < batchSize ; i ++) {
159+ int recordId = batch * batchSize + i ;
160+ ids .add ("mem_" + recordId );
161+
162+ // Create embedding
163+ List <Float > embedding = new ArrayList <>();
164+ for (int j = 0 ; j < 384 ; j ++) {
165+ embedding .add ((float ) Math .random ());
166+ }
167+ embeddings .add (embedding );
168+ }
169+
170+ collection .add ()
171+ .ids (ids )
172+ .embeddings (embeddings )
173+ .execute ();
174+
175+ // Clear local references
176+ ids = null ;
177+ embeddings = null ;
178+ }
179+
180+ long finalMemory = runtime .totalMemory () - runtime .freeMemory ();
181+ long memoryUsed = (finalMemory - initialMemory ) / (1024 * 1024 ); // MB
182+
183+ System .out .println ("Memory used: " + memoryUsed + " MB for " + totalRecords + " records" );
184+ assertTrue ("Memory usage should be reasonable" , memoryUsed < 500 ); // Less than 500MB
185+
186+ assertEquals (totalRecords , collection .count ());
187+ client .deleteCollection (collectionName );
188+ }
189+
190+ @ Test
191+ public void testQueryPerformance () throws Exception {
192+ String collectionName = "query_perf_" + UUID .randomUUID ().toString ().substring (0 , 8 );
193+ Collection collection = client .createCollection (collectionName );
194+
195+ // Add test data
196+ List <String > ids = new ArrayList <>();
197+ List <List <Float >> embeddings = new ArrayList <>();
198+ List <Map <String , Object >> metadatas = new ArrayList <>();
199+
200+ for (int i = 0 ; i < 1000 ; i ++) {
201+ ids .add ("perf_" + i );
202+
203+ List <Float > embedding = new ArrayList <>();
204+ for (int j = 0 ; j < 384 ; j ++) {
205+ embedding .add ((float ) Math .random ());
206+ }
207+ embeddings .add (embedding );
208+
209+ metadatas .add (Map .of (
210+ "type" , i % 2 == 0 ? "even" : "odd" ,
211+ "value" , i
212+ ));
213+ }
214+
215+ collection .add ()
216+ .ids (ids )
217+ .embeddings (embeddings )
218+ .metadatas (metadatas )
219+ .execute ();
220+
221+ // Measure query performance
222+ List <Float > queryEmbedding = IntStream .range (0 , 384 )
223+ .mapToObj (i -> (float ) Math .random ())
224+ .collect (Collectors .toList ());
225+
226+ long startTime = System .currentTimeMillis ();
227+
228+ for (int i = 0 ; i < 100 ; i ++) {
229+ QueryResponse result = collection .query ()
230+ .queryEmbeddings (Arrays .asList (queryEmbedding ))
231+ .nResults (10 )
232+ .where (Where .eq ("type" , "even" ))
233+ .execute ();
234+ assertNotNull (result );
235+ }
236+
237+ long duration = System .currentTimeMillis () - startTime ;
238+ double avgQueryTime = duration / 100.0 ;
239+
240+ System .out .println ("Average query time: " + avgQueryTime + " ms" );
241+ assertTrue ("Queries should be fast" , avgQueryTime < 100 ); // Less than 100ms average
242+
243+ client .deleteCollection (collectionName );
244+ }
245+ }
0 commit comments