Skip to content

Commit 0c5dd0e

Browse files
handle new chroma db api version
1 parent 037eedb commit 0c5dd0e

File tree

2 files changed

+45
-1
lines changed

2 files changed

+45
-1
lines changed

modules/chromadb/src/main/java/org/testcontainers/chromadb/ChromaDBContainer.java

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package org.testcontainers.chromadb;
22

3+
import lombok.extern.slf4j.Slf4j;
34
import org.testcontainers.containers.GenericContainer;
45
import org.testcontainers.containers.wait.strategy.Wait;
6+
import org.testcontainers.utility.ComparableVersion;
57
import org.testcontainers.utility.DockerImageName;
68

79
/**
@@ -11,6 +13,7 @@
1113
* <p>
1214
* Exposed ports: 8000
1315
*/
16+
@Slf4j
1417
public class ChromaDBContainer extends GenericContainer<ChromaDBContainer> {
1518

1619
private static final DockerImageName DEFAULT_DOCKER_IMAGE = DockerImageName.parse("chromadb/chroma");
@@ -22,13 +25,33 @@ public ChromaDBContainer(String dockerImageName) {
2225
}
2326

2427
public ChromaDBContainer(DockerImageName dockerImageName) {
28+
this(dockerImageName, isVersion2(dockerImageName.getVersionPart()));
29+
}
30+
31+
public ChromaDBContainer(DockerImageName dockerImageName, boolean isVersion2) {
2532
super(dockerImageName);
33+
34+
String apiPath = isVersion2 ? "/api/v2/heartbeat" : "/api/v1/heartbeat";
2635
dockerImageName.assertCompatibleWith(DEFAULT_DOCKER_IMAGE, GHCR_DOCKER_IMAGE);
2736
withExposedPorts(8000);
28-
waitingFor(Wait.forHttp("/api/v1/heartbeat"));
37+
waitingFor(Wait.forHttp(apiPath));
2938
}
3039

3140
public String getEndpoint() {
3241
return "http://" + getHost() + ":" + getFirstMappedPort();
3342
}
43+
44+
private static boolean isVersion2(String version) {
45+
if (version.equals("latest")) {
46+
return true;
47+
}
48+
49+
ComparableVersion comparableVersion = new ComparableVersion(version);
50+
if (comparableVersion.isGreaterThanOrEqualTo("1.0.0")) {
51+
return true;
52+
}
53+
54+
log.warn("Version {} is not a semantic version,", version);
55+
return false;
56+
}
3457
}

modules/chromadb/src/test/java/org/testcontainers/chromadb/ChromaDBContainerTest.java

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,25 @@ public void test() {
2727
given().baseUri(chroma.getEndpoint()).when().get("/api/v1/databases/test").then().statusCode(200);
2828
}
2929
}
30+
31+
@Test
32+
public void testVersion2() {
33+
try ( // container {
34+
ChromaDBContainer chroma = new ChromaDBContainer("chromadb/chroma:1.0.0")
35+
// }
36+
) {
37+
chroma.start();
38+
39+
given()
40+
.baseUri(chroma.getEndpoint())
41+
.when()
42+
.body("{\"name\": \"test\"}")
43+
.contentType(ContentType.JSON)
44+
.post("/api/v2/tenants")
45+
.then()
46+
.statusCode(200);
47+
48+
given().baseUri(chroma.getEndpoint()).when().get("/api/v2/tenants/test").then().statusCode(200);
49+
}
50+
}
3051
}

0 commit comments

Comments
 (0)