Skip to content

Commit 6a07961

Browse files
vga91vga91
committed
[NOID] test fixes
1 parent c7422be commit 6a07961

File tree

3 files changed

+58
-102
lines changed

3 files changed

+58
-102
lines changed

‎full-it/src/test/java/apoc/full/it/vectordb/ChromaDbTest.java‎

Lines changed: 15 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,6 @@
11
package apoc.full.it.vectordb;
22

3-
import apoc.ml.Prompt;
4-
import apoc.util.TestUtil;
5-
import org.junit.AfterClass;
6-
import org.junit.Before;
7-
import org.junit.BeforeClass;
8-
import org.junit.ClassRule;
9-
import org.junit.Test;
10-
import org.junit.rules.TemporaryFolder;
11-
import org.neo4j.dbms.api.DatabaseManagementService;
12-
import org.neo4j.graphdb.GraphDatabaseService;
13-
import org.neo4j.test.TestDatabaseManagementServiceBuilder;
14-
import org.testcontainers.chromadb.ChromaDBContainer;
15-
16-
import java.util.List;
17-
import java.util.Map;
18-
import java.util.concurrent.atomic.AtomicReference;
19-
20-
import static apoc.ml.RestAPIConfig.HEADERS_KEY;
3+
import static apoc.ml.Prompt.API_KEY_CONF;
214
import static apoc.util.MapUtil.map;
225
import static apoc.util.TestUtil.testCall;
236
import static apoc.util.TestUtil.testResult;
@@ -30,8 +13,6 @@
3013
import static apoc.vectordb.VectorDbTestUtil.dropAndDeleteAll;
3114
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
3215
import static apoc.vectordb.VectorDbUtil.ERROR_READONLY_MAPPING;
33-
import static apoc.vectordb.VectorDbTestUtil.getAuthHeader;
34-
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
3516
import static apoc.vectordb.VectorEmbeddingConfig.ALL_RESULTS_KEY;
3617
import static apoc.vectordb.VectorEmbeddingConfig.MAPPING_KEY;
3718
import static apoc.vectordb.VectorMappingConfig.*;
@@ -42,14 +23,14 @@
4223
import static org.neo4j.configuration.GraphDatabaseSettings.DEFAULT_DATABASE_NAME;
4324
import static org.neo4j.configuration.GraphDatabaseSettings.SYSTEM_DATABASE_NAME;
4425

26+
import apoc.ml.Prompt;
4527
import apoc.util.TestUtil;
4628
import apoc.vectordb.ChromaDb;
4729
import apoc.vectordb.VectorDb;
30+
import apoc.vectordb.VectorDbTestUtil;
4831
import java.util.List;
4932
import java.util.Map;
5033
import java.util.concurrent.atomic.AtomicReference;
51-
52-
import apoc.vectordb.VectorDbTestUtil;
5334
import org.assertj.core.api.Assertions;
5435
import org.junit.AfterClass;
5536
import org.junit.Before;
@@ -85,7 +66,7 @@ public static void setUp() throws Exception {
8566
CHROMA_CONTAINER.start();
8667

8768
HOST = CHROMA_CONTAINER.getEndpoint();
88-
TestUtil.registerProcedure(db, ChromaDb.class, VectorDb.class, Prompt.class));
69+
TestUtil.registerProcedure(db, ChromaDb.class, VectorDb.class, Prompt.class);
8970

9071
testCall(
9172
db,
@@ -437,7 +418,7 @@ public void queryVectorsWithCreateRel() {
437418
@Test
438419
public void queryVectorsWithSystemDbStorage() {
439420
String keyConfig = "chroma-config-foo";
440-
String baseUrl = "http://" + HOST;
421+
String baseUrl = HOST;
441422
Map<String, Object> mapping =
442423
map(EMBEDDING_KEY, "vect", NODE_LABEL, "Test", ENTITY_KEY, "myId", METADATA_KEY, "foo");
443424
sysDb.executeTransactionally(
@@ -480,28 +461,21 @@ public void queryVectorsWithSystemDbStorage() {
480461
public void queryVectorsWithRag() {
481462
String openAIKey = ragSetup(db);
482463

483-
Map<String, Object> conf = map(ALL_RESULTS_KEY, true,
484-
HEADERS_KEY, READONLY_AUTHORIZATION,
485-
MAPPING_KEY, map(NODE_LABEL, "Rag",
486-
ENTITY_KEY, "readID",
487-
METADATA_KEY, "foo")
488-
);
489-
490-
testResult(db,
491-
"""
492-
CALL apoc.vectordb.chroma.getAndUpdate($host, $collection, ['1', '2'], $conf) YIELD node, metadata, id, vector
493-
WITH collect(node) as paths
494-
CALL apoc.ml.rag(paths, $attributes, "Which city has foo equals to one?", $confPrompt) YIELD value
495-
RETURN value
496-
"""
497-
,
464+
Map<String, Object> conf = map(
465+
ALL_RESULTS_KEY, true, MAPPING_KEY, map(NODE_LABEL, "Rag", ENTITY_KEY, "readID", METADATA_KEY, "foo"));
466+
467+
testResult(
468+
db,
469+
"CALL apoc.vectordb.chroma.getAndUpdate($host, $collection, ['1', '2'], $conf) YIELD node, metadata, id, vector\n"
470+
+ "WITH collect(node) as paths\n"
471+
+ "CALL apoc.ml.rag(paths, $attributes, \"Which city has foo equals to one?\", $confPrompt) YIELD value\n"
472+
+ "RETURN value",
498473
map(
499474
"host", HOST,
500475
"conf", conf,
501476
"collection", COLL_ID.get(),
502477
"confPrompt", map(API_KEY_CONF, openAIKey),
503-
"attributes", List.of("city", "foo")
504-
),
478+
"attributes", List.of("city", "foo")),
505479
VectorDbTestUtil::assertRagWithVectors);
506480
}
507481
}

‎full-it/src/test/java/apoc/full/it/vectordb/QdrantTest.java‎

Lines changed: 25 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package apoc.full.it.vectordb;
22

3+
import static apoc.ml.Prompt.API_KEY_CONF;
34
import static apoc.ml.RestAPIConfig.HEADERS_KEY;
45
import static apoc.util.MapUtil.map;
56
import static apoc.util.TestUtil.testCall;
@@ -32,7 +33,6 @@
3233
import apoc.vectordb.Qdrant;
3334
import apoc.vectordb.VectorDb;
3435
import apoc.vectordb.VectorDbTestUtil;
35-
3636
import java.util.List;
3737
import java.util.Map;
3838
import org.assertj.core.api.Assertions;
@@ -210,32 +210,35 @@ public void queryVectors() {
210210

211211
@Test
212212
public void queryVectorsWithRag() {
213-
String openAIKey = System.getenv("OPENAI_KEY");;
213+
String openAIKey = System.getenv("OPENAI_KEY");
214+
;
214215
Assume.assumeNotNull("No OPENAI_KEY environment configured", openAIKey);
215216

216217
db.executeTransactionally("CREATE (:Rag {readID: 'one'}), (:Rag {readID: 'two'})");
217218

218-
Map<String, Object> conf = map(ALL_RESULTS_KEY, true,
219-
HEADERS_KEY, READONLY_AUTHORIZATION,
220-
MAPPING_KEY, map(NODE_LABEL, "Rag",
221-
ENTITY_KEY, "readID",
222-
METADATA_KEY, "foo")
223-
);
224-
225-
testResult(db,
226-
"""
227-
CALL apoc.vectordb.qdrant.getAndUpdate($host, 'test_collection', [1, 2], $conf) YIELD node, metadata, id, vector
228-
WITH collect(node) as paths
229-
CALL apoc.ml.rag(paths, $attributes, "Which city has foo equals to one?", $confPrompt) YIELD value
230-
RETURN value
231-
"""
232-
,
219+
Map<String, Object> conf = map(
220+
ALL_RESULTS_KEY,
221+
true,
222+
HEADERS_KEY,
223+
READONLY_AUTHORIZATION,
224+
MAPPING_KEY,
225+
map(NODE_LABEL, "Rag", ENTITY_KEY, "readID", METADATA_KEY, "foo"));
226+
227+
testResult(
228+
db,
229+
"CALL apoc.vectordb.qdrant.getAndUpdate($host, 'test_collection', [1, 2], $conf) YIELD node, metadata, id, vector\n"
230+
+ "WITH collect(node) as paths\n"
231+
+ "CALL apoc.ml.rag(paths, $attributes, \"Which city has foo equals to one?\", $confPrompt) YIELD value\n"
232+
+ "RETURN value",
233233
map(
234-
"host", HOST,
235-
"conf", conf,
236-
"confPrompt", map(API_KEY_CONF, openAIKey),
237-
"attributes", List.of("city", "foo")
238-
),
234+
"host",
235+
HOST,
236+
"conf",
237+
conf,
238+
"confPrompt",
239+
map(API_KEY_CONF, openAIKey),
240+
"attributes",
241+
List.of("city", "foo")),
239242
r -> {
240243
Map<String, Object> row = r.next();
241244
Object value = row.get("value");

‎full-it/src/test/java/apoc/full/it/vectordb/WeaviateTest.java‎

Lines changed: 18 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,6 @@
11
package apoc.full.it.vectordb;
22

3-
import apoc.ml.Prompt;
4-
import apoc.util.MapUtil;
5-
import apoc.util.TestUtil;
6-
import org.junit.AfterClass;
7-
import org.junit.Assume;
8-
import org.junit.Before;
9-
import org.junit.BeforeClass;
10-
import org.junit.ClassRule;
11-
import org.junit.Test;
12-
import org.junit.rules.TemporaryFolder;
13-
import org.neo4j.dbms.api.DatabaseManagementService;
14-
import org.neo4j.graphdb.Entity;
15-
import org.neo4j.graphdb.GraphDatabaseService;
16-
import org.neo4j.graphdb.ResourceIterator;
17-
import org.neo4j.test.TestDatabaseManagementServiceBuilder;
18-
import org.testcontainers.weaviate.WeaviateContainer;
19-
20-
import java.util.List;
21-
import java.util.Map;
22-
3+
import static apoc.ml.Prompt.API_KEY_CONF;
234
import static apoc.ml.RestAPIConfig.HEADERS_KEY;
245
import static apoc.util.TestUtil.testCall;
256
import static apoc.util.TestUtil.testCallEmpty;
@@ -28,18 +9,17 @@
289
import static apoc.vectordb.VectorDbHandler.Type.WEAVIATE;
2910
import static apoc.vectordb.VectorDbTestUtil.*;
3011
import static apoc.vectordb.VectorDbTestUtil.EntityType.*;
31-
import static apoc.vectordb.VectorDbUtil.ERROR_READONLY_MAPPING;
3212
import static apoc.vectordb.VectorDbTestUtil.EntityType.FALSE;
3313
import static apoc.vectordb.VectorDbTestUtil.EntityType.NODE;
3414
import static apoc.vectordb.VectorDbTestUtil.EntityType.REL;
3515
import static apoc.vectordb.VectorDbTestUtil.assertBerlinResult;
3616
import static apoc.vectordb.VectorDbTestUtil.assertLondonResult;
3717
import static apoc.vectordb.VectorDbTestUtil.assertNodesCreated;
38-
import static apoc.vectordb.VectorDbTestUtil.assertReadOnlyProcWithMappingResults;
3918
import static apoc.vectordb.VectorDbTestUtil.assertRelsCreated;
4019
import static apoc.vectordb.VectorDbTestUtil.dropAndDeleteAll;
4120
import static apoc.vectordb.VectorDbTestUtil.getAuthHeader;
4221
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
22+
import static apoc.vectordb.VectorDbUtil.ERROR_READONLY_MAPPING;
4323
import static apoc.vectordb.VectorEmbeddingConfig.ALL_RESULTS_KEY;
4424
import static apoc.vectordb.VectorEmbeddingConfig.FIELDS_KEY;
4525
import static apoc.vectordb.VectorEmbeddingConfig.MAPPING_KEY;
@@ -591,28 +571,27 @@ public void queryVectorsWithRag() {
591571
String openAIKey = ragSetup(db);
592572

593573
Map<String, Object> conf = MapUtil.map(
594-
FIELDS_KEY, FIELDS,
595-
ALL_RESULTS_KEY, true,
596-
HEADERS_KEY, READONLY_AUTHORIZATION,
597-
MAPPING_KEY, MapUtil.map(EMBEDDING_KEY, "vect",
598-
NODE_LABEL, "Rag",
599-
ENTITY_KEY, "readID",
600-
METADATA_KEY, "foo")
601-
);
602-
603-
testResult(db,
604-
"CALL apoc.vectordb.weaviate.getAndUpdate($host, 'TestCollection', [$id1], $conf) YIELD score, node, metadata, id, vector\n" +
605-
"WITH collect(node) as paths\n" +
606-
"CALL apoc.ml.rag(paths, $attributes, \"Which city has foo equals to one?\", $confPrompt) YIELD value\n" +
607-
"RETURN value"
608-
,
574+
FIELDS_KEY,
575+
FIELDS,
576+
ALL_RESULTS_KEY,
577+
true,
578+
HEADERS_KEY,
579+
READONLY_AUTHORIZATION,
580+
MAPPING_KEY,
581+
MapUtil.map(EMBEDDING_KEY, "vect", NODE_LABEL, "Rag", ENTITY_KEY, "readID", METADATA_KEY, "foo"));
582+
583+
testResult(
584+
db,
585+
"CALL apoc.vectordb.weaviate.getAndUpdate($host, 'TestCollection', [$id1], $conf) YIELD score, node, metadata, id, vector\n"
586+
+ "WITH collect(node) as paths\n"
587+
+ "CALL apoc.ml.rag(paths, $attributes, \"Which city has foo equals to one?\", $confPrompt) YIELD value\n"
588+
+ "RETURN value",
609589
MapUtil.map(
610590
"host", HOST,
611591
"id1", ID_1,
612592
"conf", conf,
613593
"confPrompt", MapUtil.map(API_KEY_CONF, openAIKey),
614-
"attributes", List.of("city", "foo")
615-
),
594+
"attributes", List.of("city", "foo")),
616595
VectorDbTestUtil::assertRagWithVectors);
617596
}
618597
}

0 commit comments

Comments
 (0)