From 13262785dedf32a97e392afc1a758616995dc9d9 Mon Sep 17 00:00:00 2001
From: Jakob Beckmann <32326425+f4z3r@users.noreply.github.com>
Date: Fri, 5 Apr 2024 17:45:41 +0200
Subject: [PATCH] fix(vault): add support for HashiCorp Vault container (#366)

Add support for a Vault container.
---
 index.rst                                     |  1 +
 modules/vault/README.rst                      |  2 +
 .../vault/testcontainers/vault/__init__.py    | 74 +++++++++++++++++++
 modules/vault/tests/test_vault.py             | 41 ++++++++++
 poetry.lock                                   | 20 ++++-
 pyproject.toml                                |  4 +
 6 files changed, 141 insertions(+), 1 deletion(-)
 create mode 100644 modules/vault/README.rst
 create mode 100644 modules/vault/testcontainers/vault/__init__.py
 create mode 100644 modules/vault/tests/test_vault.py

diff --git a/index.rst b/index.rst
index 2a2bc659..4f3dad80 100644
--- a/index.rst
+++ b/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
diff --git a/modules/vault/README.rst b/modules/vault/README.rst
new file mode 100644
index 00000000..71c079df
--- /dev/null
+++ b/modules/vault/README.rst
@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.vault.VaultContainer
+.. title:: testcontainers.vault.VaultContainer
diff --git a/modules/vault/testcontainers/vault/__init__.py b/modules/vault/testcontainers/vault/__init__.py
new file mode 100644
index 00000000..5f50cdd4
--- /dev/null
+++ b/modules/vault/testcontainers/vault/__init__.py
@@ -0,0 +1,74 @@
+#
+#    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.
+
+from http.client import HTTPException
+from urllib.request import urlopen
+
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready
+
+
+class VaultContainer(DockerContainer):
+    """
+    Vault container.
+
+    Example:
+
+        .. doctest::
+
+            >>> from testcontainers.vault import VaultContainer
+            >>> import hvac
+
+            >>> with VaultContainer("hashicorp/vault:1.16.1") as vault_container:
+            ...     connection_url = vault_container.get_connection_url()
+            ...     client = hvac.Client(url=connection_url, token=vault_container.root_token)
+            ...     assert client.is_authenticated()
+            ...     # use root client to perform desired actions, e.g.
+            ...     policies = client.sys.list_acl_policies()
+    """
+
+    def __init__(
+        self,
+        image: str = "hashicorp/vault:latest",
+        port: int = 8200,
+        root_token: str = "toor",
+        **kwargs,
+    ) -> None:
+        super().__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_connection_url(self) -> str:
+        """
+        Get the connection URL used to connect to the Vault container.
+
+        Returns:
+            str: The address to connect to.
+        """
+        host_ip = self.get_container_host_ip()
+        exposed_port = self.get_exposed_port(self.port)
+        return f"http://{host_ip}:{exposed_port}"
+
+    @wait_container_is_ready(HTTPException)
+    def _healthcheck(self) -> None:
+        url = f"{self.get_connection_url()}/v1/sys/health"
+        with urlopen(url) as res:
+            if res.status > 299:
+                raise HTTPException()
+
+    def start(self) -> "VaultContainer":
+        super().start()
+        self._healthcheck()
+        return self
diff --git a/modules/vault/tests/test_vault.py b/modules/vault/tests/test_vault.py
new file mode 100644
index 00000000..54017d2f
--- /dev/null
+++ b/modules/vault/tests/test_vault.py
@@ -0,0 +1,41 @@
+import hvac
+from testcontainers.vault import VaultContainer
+
+
+def test_docker_run_vault():
+    config = VaultContainer("hashicorp/vault:1.16.1")
+    with config as vault:
+        url = vault.get_connection_url()
+        client = hvac.Client(url=url)
+        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:
+        url = vault.get_connection_url()
+        client = hvac.Client(url=url, token=vault.root_token)
+        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"
diff --git a/poetry.lock b/poetry.lock
index 614e42e5..a2f81d3b 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1449,6 +1449,23 @@ cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"]
 http2 = ["h2 (>=3,<5)"]
 socks = ["socksio (==1.*)"]
 
+[[package]]
+name = "hvac"
+version = "2.1.0"
+description = "HashiCorp Vault API client"
+optional = false
+python-versions = ">=3.8,<4.0"
+files = [
+    {file = "hvac-2.1.0-py3-none-any.whl", hash = "sha256:73bc91e58c3fc7c6b8107cdaca9cb71fa0a893dfd80ffbc1c14e20f24c0c29d7"},
+    {file = "hvac-2.1.0.tar.gz", hash = "sha256:b48bcda11a4ab0a7b6c47232c7ba7c87fda318ae2d4a7662800c465a78742894"},
+]
+
+[package.dependencies]
+requests = ">=2.27.1,<3.0.0"
+
+[package.extras]
+parser = ["pyhcl (>=0.4.4,<0.5.0)"]
+
 [[package]]
 name = "hyperframe"
 version = "6.0.1"
@@ -4188,9 +4205,10 @@ rabbitmq = ["pika"]
 redis = ["redis"]
 registry = ["bcrypt"]
 selenium = ["selenium"]
+vault = []
 weaviate = ["weaviate-client"]
 
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9,<4.0"
-content-hash = "af9f21cb52ebd761ba91c852c9839d982124c149e77abdc079fc5657cecb9ff7"
+content-hash = "233dfd72d07a555973aafc3fe3b6676574403b9fe4bb2c0230d455cff8aa2933"
diff --git a/pyproject.toml b/pyproject.toml
index 2d2fbeb5..08c9c68b 100644
--- a/pyproject.toml
+++ b/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" }
 ]
 
@@ -125,6 +126,7 @@ rabbitmq = ["pika"]
 redis = ["redis"]
 registry = ["bcrypt"]
 selenium = ["selenium"]
+vault = []
 weaviate = ["weaviate-client"]
 chroma = ["chromadb-client"]
 
@@ -144,6 +146,7 @@ psycopg = "*"
 cassandra-driver = "*"
 pytest-asyncio = "0.23.5"
 kafka-python-ng = "^2.2.0"
+hvac = "*"
 
 [[tool.poetry.source]]
 name = "PyPI"
@@ -262,6 +265,7 @@ mypy_path = [
 #    "modules/rabbitmq",
 #    "modules/redis",
 #    "modules/selenium"
+#    "modules/vault"
 #    "modules/weaviate"
 ]
 enable_error_code = [