11package org .testcontainers .chromadb ;
22
3+ import lombok .extern .slf4j .Slf4j ;
34import org .testcontainers .containers .GenericContainer ;
45import org .testcontainers .containers .wait .strategy .Wait ;
6+ import org .testcontainers .utility .ComparableVersion ;
57import org .testcontainers .utility .DockerImageName ;
68
79/**
1113 * <p>
1214 * Exposed ports: 8000
1315 */
16+ @ Slf4j
1417public class ChromaDBContainer extends GenericContainer <ChromaDBContainer > {
1518
1619 private static final DockerImageName DEFAULT_DOCKER_IMAGE = DockerImageName .parse ("chromadb/chroma" );
@@ -22,13 +25,32 @@ 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+ String apiPath = isVersion2 ? "/api/v2/heartbeat" : "/api/v1/heartbeat" ;
2634 dockerImageName .assertCompatibleWith (DEFAULT_DOCKER_IMAGE , GHCR_DOCKER_IMAGE );
2735 withExposedPorts (8000 );
28- waitingFor (Wait .forHttp ("/api/v1/heartbeat" ));
36+ waitingFor (Wait .forHttp (apiPath ));
2937 }
3038
3139 public String getEndpoint () {
3240 return "http://" + getHost () + ":" + getFirstMappedPort ();
3341 }
42+
43+ private static boolean isVersion2 (String version ) {
44+ if (version .equals ("latest" )) {
45+ return true ;
46+ }
47+
48+ ComparableVersion comparableVersion = new ComparableVersion (version );
49+ if (comparableVersion .isGreaterThanOrEqualTo ("1.0.0" )) {
50+ return true ;
51+ }
52+
53+ log .warn ("Version {} is less than 1.0.0 or not a semantic version." , version );
54+ return false ;
55+ }
3456}
0 commit comments