Skip to content

Commit 90762e8

Browse files
antas-marcinalexanderankin
antas-marcin
and
alexanderankin
authoredMar 30, 2024··
fix: Add Weaviate module (#492)
This PR adds Weaviate module. --------- Co-authored-by: David Ankin <daveankin@gmail.com>
·
testcontainers-v4.9.2testcontainers-v4.3.0
1 parent 3d891a5 commit 90762e8

File tree

6 files changed

+549
-2
lines changed

6 files changed

+549
-2
lines changed
 

‎index.rst

+1
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
3737
modules/rabbitmq/README
3838
modules/redis/README
3939
modules/selenium/README
40+
modules/weaviate/README
4041

4142
Getting Started
4243
---------------

‎modules/weaviate/README.rst

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
.. autoclass:: testcontainers.weaviate.WeaviateContainer
2+
.. title:: testcontainers.weaviate.WeaviateContainer

‎modules/weaviate/testcontainers/weaviate/__init__.py

+178
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
#
2+
# Licensed under the Apache License, Version 2.0 (the "License"); you may
3+
# not use this file except in compliance with the License. You may obtain
4+
# a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
10+
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
11+
# License for the specific language governing permissions and limitations
12+
# under the License.
13+
from typing import TYPE_CHECKING, Optional
14+
15+
from requests import ConnectionError, get
16+
17+
from testcontainers.core.generic import DbContainer
18+
from testcontainers.core.waiting_utils import wait_container_is_ready
19+
20+
if TYPE_CHECKING:
21+
from requests import Response
22+
23+
24+
class WeaviateContainer(DbContainer):
25+
"""
26+
Weaviate vector database container.
27+
28+
Arguments:
29+
`image`
30+
Docker image to use with Weaviate container.
31+
`env_vars`
32+
Additional environment variables to include with the container, e.g. ENABLE_MODULES list, QUERY_DEFAULTS_LIMIT setting.
33+
34+
Example:
35+
This example shows how to start Weaviate container with defualt settings.
36+
37+
.. doctest::
38+
39+
>>> from testcontainers.weaviate import WeaviateContainer
40+
41+
>>> with WeaviateContainer() as container:
42+
... with container.get_client() as client:
43+
... client.is_live()
44+
True
45+
46+
This example shows how to start Weaviate container with additinal settings.
47+
48+
.. doctest::
49+
50+
>>> from testcontainers.weaviate import WeaviateContainer
51+
52+
>>> with WeaviateContainer(
53+
... env_vars={
54+
... "ENABLE_MODULES": "backup-filesystem,text2vec-openai",
55+
... "BACKUP_FILESYSTEM_PATH": "/tmp/backups",
56+
... "QUERY_DEFAULTS_LIMIT": 100,
57+
... }
58+
... ) as container:
59+
... with container.get_client() as client:
60+
... client.is_live()
61+
True
62+
"""
63+
64+
def __init__(
65+
self,
66+
image: str = "semitechnologies/weaviate:1.24.5",
67+
env_vars: Optional[dict[str, str]] = None,
68+
**kwargs,
69+
) -> None:
70+
super().__init__(image, **kwargs)
71+
self._http_port = 8080
72+
self._grpc_port = 50051
73+
74+
self.with_command(f"--host 0.0.0.0 --scheme http --port {self._http_port}")
75+
self.with_exposed_ports(self._http_port, self._grpc_port)
76+
77+
if env_vars is not None:
78+
for key, value in env_vars.items():
79+
self.with_env(key, value)
80+
81+
def _configure(self) -> None:
82+
self.with_env("AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED", "true")
83+
self.with_env("PERSISTENCE_DATA_PATH", "/var/lib/weaviate")
84+
85+
@wait_container_is_ready(ConnectionError)
86+
def _connect(self) -> None:
87+
url = f"http://{self.get_http_host()}:{self.get_http_port()}/v1/.well-known/ready"
88+
response: Response = get(url)
89+
response.raise_for_status()
90+
91+
def get_client(
92+
self,
93+
headers: Optional[dict[str, str]] = None,
94+
):
95+
"""
96+
Get a `weaviate.WeaviateClient` instance associated with the container.
97+
98+
Arguments:
99+
`headers`
100+
Additional headers to include in the requests, e.g. API keys for third-party Cloud vectorization.
101+
102+
Returns:
103+
WeaviateClient: An instance of the `weaviate.WeaviateClient` class.
104+
"""
105+
106+
try:
107+
import weaviate
108+
except ImportError as e:
109+
raise ImportError("To use the `get_client` method, you must install the `weaviate-client` package.") from e
110+
return weaviate.connect_to_custom(
111+
http_host=self.get_http_host(),
112+
http_port=self.get_http_port(),
113+
http_secure=self.get_http_secure(),
114+
grpc_host=self.get_http_host(),
115+
grpc_port=self.get_grpc_port(),
116+
grpc_secure=self.get_grpc_secure(),
117+
headers=headers,
118+
)
119+
120+
def get_http_host(self) -> str:
121+
"""
122+
Get the HTTP host of Weaviate container.
123+
124+
Returns:
125+
`str`
126+
The HTTP host of Weaviate container.
127+
"""
128+
return f"{self.get_container_host_ip()}"
129+
130+
def get_http_port(self) -> int:
131+
"""
132+
Get the HTTP port of Weaviate container.
133+
134+
Returns:
135+
`int`
136+
The HTTP port of Weaviate container.
137+
"""
138+
return self.get_exposed_port(self._http_port)
139+
140+
def get_http_secure(self) -> bool:
141+
"""
142+
Get the HTTP secured setting of Weaviate container.
143+
144+
Returns:
145+
`bool`
146+
True if it's https.
147+
"""
148+
return False
149+
150+
def get_grpc_host(self) -> str:
151+
"""
152+
Get the gRPC host of Weaviate container.
153+
154+
Returns:
155+
`str`
156+
The gRPC host of Weaviate container.
157+
"""
158+
return f"{self.get_container_host_ip()}"
159+
160+
def get_grpc_port(self) -> int:
161+
"""
162+
Get the gRPC port of Weaviate container.
163+
164+
Returns:
165+
`int`
166+
The gRPC port of Weaviate container.
167+
"""
168+
return self.get_exposed_port(self._grpc_port)
169+
170+
def get_grpc_secure(self) -> bool:
171+
"""
172+
Get the gRPC secured setting of Weaviate container.
173+
174+
Returns:
175+
`str`
176+
True if the conntection is secured with SSL.
177+
"""
178+
return False

‎modules/weaviate/tests/test_weaviate.py

+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
from testcontainers.weaviate import WeaviateContainer
2+
import weaviate
3+
4+
5+
def test_docker_run_weaviate():
6+
with WeaviateContainer() as container:
7+
client = weaviate.connect_to_custom(
8+
http_host=container.get_http_host(),
9+
http_port=container.get_http_port(),
10+
http_secure=container.get_http_secure(),
11+
grpc_host=container.get_grpc_host(),
12+
grpc_port=container.get_grpc_port(),
13+
grpc_secure=container.get_grpc_secure(),
14+
)
15+
16+
meta = client.get_meta()
17+
assert len(meta.get("version")) > 0
18+
19+
client.close()
20+
21+
22+
def test_docker_run_weaviate_with_client():
23+
with WeaviateContainer() as container:
24+
with container.get_client() as client:
25+
assert client.is_live()
26+
27+
meta = client.get_meta()
28+
assert len(meta.get("version")) > 0
29+
30+
31+
def test_docker_run_weaviate_with_modules():
32+
enable_modules = [
33+
"backup-filesystem",
34+
"text2vec-openai",
35+
"text2vec-cohere",
36+
"text2vec-huggingface",
37+
"generative-openai",
38+
]
39+
with WeaviateContainer(
40+
env_vars={
41+
"ENABLE_MODULES": ",".join(enable_modules),
42+
"BACKUP_FILESYSTEM_PATH": "/tmp/backups",
43+
}
44+
) as container:
45+
with container.get_client() as client:
46+
assert client.is_live()
47+
48+
meta = client.get_meta()
49+
assert len(meta.get("version")) > 0
50+
51+
modules = meta.get("modules")
52+
assert len(modules) == len(enable_modules)
53+
54+
for name in enable_modules:
55+
assert len(modules[name]) > 0

‎poetry.lock

+308-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml

+5-1
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,8 @@ packages = [
4949
{ include = "testcontainers", from = "modules/postgres" },
5050
{ include = "testcontainers", from = "modules/rabbitmq" },
5151
{ include = "testcontainers", from = "modules/redis" },
52-
{ include = "testcontainers", from = "modules/selenium" }
52+
{ include = "testcontainers", from = "modules/selenium" },
53+
{ include = "testcontainers", from = "modules/weaviate" }
5354
]
5455

5556
[tool.poetry.urls]
@@ -85,6 +86,7 @@ cx_Oracle = { version = "*", optional = true }
8586
pika = { version = "*", optional = true }
8687
redis = { version = "*", optional = true }
8788
selenium = { version = "*", optional = true }
89+
weaviate-client = { version = "^4.5.4", optional = true }
8890

8991
[tool.poetry.extras]
9092
arangodb = ["python-arango"]
@@ -109,6 +111,7 @@ postgres = []
109111
rabbitmq = ["pika"]
110112
redis = ["redis"]
111113
selenium = ["selenium"]
114+
weaviate = ["weaviate-client"]
112115

113116
[tool.poetry.group.dev.dependencies]
114117
mypy = "1.7.1"
@@ -241,6 +244,7 @@ mypy_path = [
241244
# "modules/rabbitmq",
242245
# "modules/redis",
243246
# "modules/selenium"
247+
# "modules/weaviate"
244248
]
245249
enable_error_code = [
246250
"ignore-without-code",

0 commit comments

Comments
 (0)
Please sign in to comment.