Skip to content

Commit f6c8288

Browse files
committed
more test
1 parent c575a1f commit f6c8288

25 files changed

+705
-367
lines changed

astra-db-java/src/test/java/com/datastax/astra/test/integration/AbstractDataAPITest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public abstract class AbstractDataAPITest {
3434
public static final String ENV_VAR_DESTINATION_ASTRA_PROD = "astra_prod";
3535
public static final String ENV_VAR_DESTINATION_ASTRA_TEST = "astra_test";
3636

37-
private String readEnvVariable(String key) {
37+
public String readEnvVariable(String key) {
3838
String envVar = System.getenv(key);
3939
String systemEnvVar = System.getProperty(key);
4040
if (Utils.hasLength(envVar)) {

astra-db-java/src/test/java/com/datastax/astra/test/integration/AbstractTableITTest.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import org.junit.jupiter.api.MethodOrderer;
55
import org.junit.jupiter.api.TestMethodOrder;
66

7+
import java.util.Map;
8+
79
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
810

911
@Slf4j
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
package com.datastax.astra.test.integration;
2+
3+
import com.datastax.astra.client.collections.Collection;
4+
import com.datastax.astra.client.collections.commands.options.CollectionFindOneOptions;
5+
import com.datastax.astra.client.collections.commands.options.CollectionInsertManyOptions;
6+
import com.datastax.astra.client.collections.commands.results.CollectionInsertManyResult;
7+
import com.datastax.astra.client.collections.definition.CollectionDefinition;
8+
import com.datastax.astra.client.collections.definition.documents.Document;
9+
import com.datastax.astra.client.core.DataAPIKeywords;
10+
import com.datastax.astra.client.core.headers.EmbeddingAPIKeyHeaderProvider;
11+
import com.datastax.astra.client.core.query.Projection;
12+
import com.datastax.astra.client.core.query.Sort;
13+
import com.datastax.astra.client.core.vector.SimilarityMetric;
14+
import com.datastax.astra.client.core.vectorize.EmbeddingProvider;
15+
import com.datastax.astra.test.model.TestDataset;
16+
import lombok.extern.slf4j.Slf4j;
17+
import org.junit.jupiter.api.Test;
18+
19+
import java.util.Optional;
20+
21+
import static org.assertj.core.api.Assertions.assertThat;
22+
23+
@Slf4j
24+
public abstract class AbstractVectorizeApiHeaderITTest extends AbstractVectorizeITTest {
25+
26+
/**
27+
* Find the Embedding Provider if it exists.
28+
*/
29+
@Test
30+
public void should_test_vectorize_with_api_header() {
31+
getDatabase()
32+
.getDatabaseAdmin()
33+
.findEmbeddingProviders()
34+
.getEmbeddingProviders()
35+
.computeIfPresent(getEmbeddingProviderId(), (key, value) -> {
36+
this.testVectorizeWithApiKeyHeader(value);
37+
return value;
38+
});
39+
}
40+
41+
/**
42+
* Coming from findEmbeddingProviders, test the provider.
43+
* @param provider
44+
* provider information
45+
*/
46+
public void testVectorizeWithApiKeyHeader(EmbeddingProvider provider) {
47+
log.info("Testing embedding provider {}", getEmbeddingProviderId());
48+
49+
// Header authentication
50+
final EmbeddingAPIKeyHeaderProvider authProvider = (getApiKey() != null) ?
51+
new EmbeddingAPIKeyHeaderProvider(getApiKey()) : null;
52+
53+
// Test for each model
54+
provider.getModels().forEach(model -> {
55+
try {
56+
log.info("Testing model {}", model);
57+
58+
// (1) Create a collection with a name matching the model
59+
CollectionDefinition definition = new CollectionDefinition()
60+
.vectorSimilarity(SimilarityMetric.COSINE)
61+
.vectorize(getEmbeddingProviderId(), model.getName(), null, getAuthenticationParameters());
62+
if (model.getVectorDimension() != null) {
63+
definition.vectorDimension(model.getVectorDimension());
64+
}
65+
String collectionName = getCollectionNameFromModel(model.getName());
66+
Collection<Document> collection = getDatabase().createCollection(collectionName, definition);
67+
log.info("Collection created {}", collection.getCollectionName());
68+
69+
// (2) Ingestion
70+
CollectionInsertManyOptions options = new CollectionInsertManyOptions();
71+
if (authProvider != null) {
72+
options.embeddingAuthProvider(authProvider);
73+
}
74+
CollectionInsertManyResult res = collection.insertMany(TestDataset.DOCS_SONG_DIRE_STRAITS, options);
75+
assertThat(res.getInsertedIds()).hasSize(8);
76+
log.info("{} Documents inserted", res.getInsertedIds().size());
77+
78+
// (3) Find with Vectorize
79+
Optional<Document> doc = collection.findOne(null,
80+
new CollectionFindOneOptions()
81+
.sort(Sort.vectorize("You shouldn't come around here singing up at people like tha"))
82+
.projection(Projection.exclude(DataAPIKeywords.VECTOR.getKeyword()))
83+
.embeddingAuthProvider(authProvider)
84+
.includeSimilarity(true));
85+
log.info("Document found {}", doc);
86+
assertThat(doc).isPresent();
87+
assertThat(doc.get().getId(Integer.class)).isEqualTo(7);
88+
assertThat(doc.get().getDouble(DataAPIKeywords.SIMILARITY.getKeyword())).isGreaterThan(.8);
89+
90+
collection.drop();
91+
} catch(Exception e) {
92+
log.error("Error while testing model {}", model, e);
93+
}
94+
});
95+
}
96+
97+
98+
}

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy