|
10 | 10 | # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
11 | 11 | # License for the specific language governing permissions and limitations |
12 | 12 | # under the License. |
13 | | -from typing import TYPE_CHECKING, Optional |
| 13 | +import contextlib |
| 14 | +from typing import Optional |
14 | 15 |
|
15 | | -from requests import ConnectionError, get |
16 | | -from testcontainers.core.generic import DbContainer |
17 | | -from testcontainers.core.waiting_utils import wait_container_is_ready |
| 16 | +from docker.types import IPAMConfig, IPAMPool |
| 17 | +from testcontainers.core.container import DockerContainer |
| 18 | +from testcontainers.core.network import Network |
| 19 | +from testcontainers.core.wait_strategies import HttpWaitStrategy |
| 20 | +from typing_extensions import Self |
18 | 21 |
|
19 | | -if TYPE_CHECKING: |
20 | | - from requests import Response |
21 | 22 |
|
22 | | - |
23 | | -class WeaviateContainer(DbContainer): |
| 23 | +class WeaviateContainer(DockerContainer): |
24 | 24 | """ |
25 | 25 | Weaviate vector database container. |
26 | 26 |
|
@@ -62,30 +62,46 @@ class WeaviateContainer(DbContainer): |
62 | 62 |
|
63 | 63 | def __init__( |
64 | 64 | self, |
65 | | - image: str = "semitechnologies/weaviate:1.24.5", |
| 65 | + image: str = "semitechnologies/weaviate:1.28.4", |
66 | 66 | env_vars: Optional[dict[str, str]] = None, |
67 | 67 | **kwargs, |
68 | 68 | ) -> None: |
69 | 69 | super().__init__(image, **kwargs) |
70 | 70 | self._http_port = 8080 |
71 | 71 | self._grpc_port = 50051 |
72 | 72 |
|
73 | | - self.with_command(f"--host 0.0.0.0 --scheme http --port {self._http_port}") |
| 73 | + # Weaviate's memberlist gossip requires an RFC1918 private IP to advertise. |
| 74 | + # Docker bridge subnets are not always in RFC1918 ranges, so we create a |
| 75 | + # dedicated network using 10.10.10.0/24 which is always recognised as private. |
| 76 | + self._weaviate_network = Network( |
| 77 | + docker_network_kw={"ipam": IPAMConfig(pool_configs=[IPAMPool(subnet="10.10.10.0/24")])} |
| 78 | + ) |
| 79 | + self.with_network(self._weaviate_network) |
74 | 80 | self.with_exposed_ports(self._http_port, self._grpc_port) |
| 81 | + self.waiting_for(HttpWaitStrategy(self._http_port, "/v1/.well-known/ready")) |
75 | 82 |
|
76 | 83 | if env_vars is not None: |
77 | 84 | for key, value in env_vars.items(): |
78 | 85 | self.with_env(key, value) |
79 | 86 |
|
80 | 87 | def _configure(self) -> None: |
| 88 | + self.with_env("HOST", "0.0.0.0") |
| 89 | + self.with_env("PORT", str(self._http_port)) |
| 90 | + self.with_env("SCHEME", "http") |
81 | 91 | self.with_env("AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED", "true") |
82 | 92 | self.with_env("PERSISTENCE_DATA_PATH", "/var/lib/weaviate") |
83 | | - |
84 | | - @wait_container_is_ready(ConnectionError) |
85 | | - def _connect(self) -> None: |
86 | | - url = f"http://{self.get_http_host()}:{self.get_http_port()}/v1/.well-known/ready" |
87 | | - response: Response = get(url) |
88 | | - response.raise_for_status() |
| 93 | + self.with_env("DEFAULT_VECTORIZER_MODULE", "none") |
| 94 | + self.with_env("CLUSTER_HOSTNAME", "node1") |
| 95 | + |
| 96 | + def start(self) -> Self: |
| 97 | + self._weaviate_network.create() |
| 98 | + return super().start() |
| 99 | + |
| 100 | + def stop(self, force: bool = True, delete_volume: bool = True) -> None: |
| 101 | + super().stop(force=force, delete_volume=delete_volume) |
| 102 | + if self._weaviate_network._network is not None: |
| 103 | + with contextlib.suppress(Exception): |
| 104 | + self._weaviate_network.remove() |
89 | 105 |
|
90 | 106 | def get_client( |
91 | 107 | self, |
|
0 commit comments