Skip to content

Commit

Permalink
[qdrant] Allow to set key and config file (#8556)
Browse files Browse the repository at this point in the history
Co-authored-by: Anush008 <anushshetty90@gmail.com>
  • Loading branch information
eddumelendez and Anush008 authored Apr 18, 2024
1 parent 8d0cf9b commit ba25e5d
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.images.builder.Transferable;
import org.testcontainers.utility.DockerImageName;

/**
Expand All @@ -12,25 +13,45 @@
* Exposed ports:
* <ul>
* <li>HTTP: 6333</li>
* <li>Grpc: 6334</li>
* <li>GRPC: 6334</li>
* </ul>
*/
public class QdrantContainer extends GenericContainer<QdrantContainer> {

private static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("qdrant/qdrant");

private static final int QDRANT_REST_PORT = 6333;

private static final int QDRANT_GRPC_PORT = 6334;

private static final String CONFIG_FILE_PATH = "/qdrant/config/config.yaml";

private static final String API_KEY_ENV = "QDRANT__SERVICE__API_KEY";

public QdrantContainer(String image) {
this(DockerImageName.parse(image));
}

public QdrantContainer(DockerImageName dockerImageName) {
super(dockerImageName);
dockerImageName.assertCompatibleWith(DEFAULT_IMAGE_NAME);
withExposedPorts(6333, 6334);
waitingFor(Wait.forHttp("/readyz").forPort(6333));
withExposedPorts(QDRANT_REST_PORT, QDRANT_GRPC_PORT);
waitingFor(Wait.forHttp("/readyz").forPort(QDRANT_REST_PORT));
}

public QdrantContainer withApiKey(String apiKey) {
return withEnv(API_KEY_ENV, apiKey);
}

public QdrantContainer withConfigFile(Transferable configFile) {
return withCopyToContainer(configFile, CONFIG_FILE_PATH);
}

public int getGrpcPort() {
return getMappedPort(QDRANT_GRPC_PORT);
}

public String getGrpcHostAddress() {
return getHost() + ":" + getMappedPort(6334);
return getHost() + ":" + getGrpcPort();
}
}
Original file line number Diff line number Diff line change
@@ -1,36 +1,89 @@
package org.testcontainers.qdrant;

import io.grpc.Grpc;
import io.grpc.InsecureChannelCredentials;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.QdrantOuterClass;
import org.junit.Test;
import org.testcontainers.images.builder.Transferable;

import java.util.UUID;
import java.util.concurrent.ExecutionException;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

public class QdrantContainerTest {

@Test
public void test() throws ExecutionException, InterruptedException {
public void shouldReturnVersion() throws ExecutionException, InterruptedException {
try (
// qdrantContainer {
QdrantContainer qdrant = new QdrantContainer("qdrant/qdrant:v1.7.4")
// }
) {
qdrant.start();

QdrantGrpcClient client = QdrantGrpcClient
.newBuilder(
Grpc.newChannelBuilder(qdrant.getGrpcHostAddress(), InsecureChannelCredentials.create()).build()
)
.build();
QdrantOuterClass.HealthCheckReply healthCheckReply = client
.qdrant()
.healthCheck(QdrantOuterClass.HealthCheckRequest.getDefaultInstance())
.get();
QdrantClient client = new QdrantClient(
QdrantGrpcClient.newBuilder(qdrant.getHost(), qdrant.getGrpcPort(), false).build()
);
QdrantOuterClass.HealthCheckReply healthCheckReply = client.healthCheckAsync().get();
assertThat(healthCheckReply.getVersion()).isEqualTo("1.7.4");

client.close();
}
}

@Test
public void shouldSetApiKey() throws ExecutionException, InterruptedException {
String apiKey = UUID.randomUUID().toString();
try (QdrantContainer qdrant = new QdrantContainer("qdrant/qdrant:v1.7.4").withApiKey(apiKey)) {
qdrant.start();

final QdrantClient unauthClient = new QdrantClient(
QdrantGrpcClient.newBuilder(qdrant.getHost(), qdrant.getGrpcPort(), false).build()
);

assertThatThrownBy(() -> unauthClient.healthCheckAsync().get()).isInstanceOf(ExecutionException.class);

unauthClient.close();

final QdrantClient client = new QdrantClient(
QdrantGrpcClient.newBuilder(qdrant.getHost(), qdrant.getGrpcPort(), false).withApiKey(apiKey).build()
);

QdrantOuterClass.HealthCheckReply healthCheckReply = client.healthCheckAsync().get();
assertThat(healthCheckReply.getVersion()).isEqualTo("1.7.4");

client.close();
}
}

@Test
public void shouldSetApiKeyUsingConfigFile() throws ExecutionException, InterruptedException {
String apiKey = UUID.randomUUID().toString();
String configFile = "service:\n api_key: " + apiKey;
try (
QdrantContainer qdrant = new QdrantContainer("qdrant/qdrant:v1.7.4")
.withConfigFile(Transferable.of(configFile))
) {
qdrant.start();

final QdrantClient unauthClient = new QdrantClient(
QdrantGrpcClient.newBuilder(qdrant.getHost(), qdrant.getGrpcPort(), false).build()
);

assertThatThrownBy(() -> unauthClient.healthCheckAsync().get()).isInstanceOf(ExecutionException.class);

unauthClient.close();

final QdrantClient client = new QdrantClient(
QdrantGrpcClient.newBuilder(qdrant.getHost(), qdrant.getGrpcPort(), false).withApiKey(apiKey).build()
);

QdrantOuterClass.HealthCheckReply healthCheckReply = client.healthCheckAsync().get();
assertThat(healthCheckReply.getVersion()).isEqualTo("1.7.4");

client.close();
}
}
}

0 comments on commit ba25e5d

Please sign in to comment.