fix(vault): add support for HashiCorp Vault container

index.rst

@@ -42,6 +42,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
     modules/redis/README
     modules/registry/README
     modules/selenium/README
+    modules/vault/README
     modules/weaviate/README
 
 Getting Started

modules/vault/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.vault.VaultContainer
+.. title:: testcontainers.vault.VaultContainer

modules/vault/testcontainers/vault/__init__.py

@@ -0,0 +1,87 @@
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import hvac
+import requests
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.utils import raise_for_deprecated_parameter
+from testcontainers.core.waiting_utils import wait_container_is_ready
+
+
+class VaultContainer(DockerContainer):
+    """
+    Vault container.
+
+    Example:
+
+        .. doctest::
+
+            >>> from testcontainers.vault import VaultContainer
+
+            >>> with VaultContainer() as vault_container:
+            ...     vault_client = vault_container.get_client()
+    """
+    def __init__(self, image: str = "hashicorp/vault:latest", port: int = 8200,
+                 root_token: str = "toor", **kwargs) -> None:
+        raise_for_deprecated_parameter(kwargs, "port_to_expose", "port")
+        super(VaultContainer, self).__init__(image, **kwargs)
+        self.port = port
+        self.root_token = root_token
+        self.with_exposed_ports(self.port)
+        self.with_env("VAULT_DEV_ROOT_TOKEN_ID", self.root_token)
+
+    def get_config(self) -> dict:
+        """
+        Get the configuration used to connect to the Vault container, including the address to
+        connect to, and the root token.
+
+        Returns:
+            dict: {`address`: str, `root_token`: str}
+        """
+        host_ip = self.get_container_host_ip()
+        exposed_port = self.get_exposed_port(self.port)
+        return {
+            "root_token": self.root_token,
+            "address": f"http://{host_ip}:{exposed_port}",
+        }
+
+    def get_client(self) -> hvac.Client:
+        """
+        Get a Vault client.
+
+        Returns:
+            client: Vault client to connect to the container.
+        """
+        config = self.get_config()
+        return hvac.Client(url=config["address"])
+
+    def get_root_client(self) -> hvac.Client:
+        """
+        Get an authenticated Vault client with root token.
+
+        Returns:
+            client: Vault client to connect to the container.
+        """
+        config = self.get_config()
+        return hvac.Client(url=config["address"], token=config["root_token"])
+
+    @wait_container_is_ready(requests.ConnectionError)
+    def _healthcheck(self) -> None:
+        url = f"{self.get_config()['address']}/v1/sys/health"
+        response = requests.get(url, timeout=3)
+        response.raise_for_status()
+
+    def start(self) -> "VaultContainer":
+        super().start()
+        self._healthcheck()
+        return self

modules/vault/tests/test_vault.py

@@ -0,0 +1,39 @@
+from testcontainers.vault import VaultContainer
+
+
+def test_docker_run_vault():
+    config = VaultContainer("hashicorp/vault:1.16.1")
+    with config as vault:
+        client = vault.get_client()
+        assert not client.is_authenticated()
+        status = client.sys.read_health_status()
+        assert status.status_code == 200
+
+
+def test_docker_run_vault_act_as_root():
+    config = VaultContainer("hashicorp/vault:1.16.1")
+    with config as vault:
+        client = vault.get_root_client()
+        assert client.is_authenticated()
+        assert client.sys.is_initialized()
+        assert not client.sys.is_sealed()
+
+        client.sys.enable_secrets_engine(
+            backend_type="kv",
+            path="secrets",
+            config={
+                "version": "2",
+            },
+        )
+        client.secrets.kv.v2.create_or_update_secret(
+            path="my-secret",
+            mount_point="secrets",
+            secret={
+                "pssst": "this is secret",
+            },
+        )
+        resp = client.secrets.kv.v2.read_secret(
+            path="my-secret",
+            mount_point="secrets",
+        )
+        assert resp["data"]["data"]["pssst"] == "this is secret"

pyproject.toml

@@ -55,6 +55,7 @@ packages = [
     { include = "testcontainers", from = "modules/redis" },
     { include = "testcontainers", from = "modules/registry" },
     { include = "testcontainers", from = "modules/selenium" },
+    { include = "testcontainers", from = "modules/vault" },
     { include = "testcontainers", from = "modules/weaviate" }
 ]
 
@@ -74,6 +75,7 @@ azure-storage-blob = { version = "^12.19", optional = true }
 clickhouse-driver = { version = "*", optional = true }
 google-cloud-pubsub = { version = ">=2", optional = true }
 google-cloud-datastore = { version = ">=2", optional = true }
+hvac = { version = ">=2", optional = true }
 influxdb = { version = "*", optional = true }
 influxdb-client = { version = "*", optional = true }
 kubernetes = { version = "*", optional = true }
@@ -91,6 +93,7 @@ opensearch-py = { version = "*", optional = true }
 oracledb = { version = "*", optional = true }
 pika = { version = "*", optional = true }
 redis = { version = "*", optional = true }
+requests = { version = ">=2", optional = true }
 selenium = { version = "*", optional = true }
 weaviate-client = { version = "^4.5.4", optional = true }
 chromadb-client = { version = "*", optional = true }
@@ -125,6 +128,7 @@ rabbitmq = ["pika"]
 redis = ["redis"]
 registry = ["bcrypt"]
 selenium = ["selenium"]
+vault = ["hvac", "requests"]
 weaviate = ["weaviate-client"]
 chroma = ["chromadb-client"]
 
@@ -262,6 +266,7 @@ mypy_path = [
 #    "modules/rabbitmq",
 #    "modules/redis",
 #    "modules/selenium"
+#    "modules/vault"
 #    "modules/weaviate"
 ]
 enable_error_code = [