Add CassandraContainer

Author: cbornet

.github/settings.yml

@@ -63,6 +63,7 @@ labels:
   - { name: '🐧 linux',                  color: '#3ED4D',, description: '' }
   - { name: '👀 requires attention',     color: '#fef2c0', description: '' }
   - { name: '📖 documentation',          color: '#d93f0b', description: '' }
+  - { name: '📦 package: cassandra',     color: '#0052CC', description: '' }
   - { name: '📦 package: clickhouse',    color: '#0052CC', description: '' }
   - { name: '📦 package: compose',       color: '#0052CC', description: '' }
   - { name: '📦 package: core',          color: '#0052CC', description: '' }

INDEX.rst

@@ -17,6 +17,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
     core/README
     modules/arangodb/README
     modules/azurite/README
+    modules/cassandra/README
     modules/clickhouse/README
     modules/elasticsearch/README
     modules/google/README

modules/cassandra/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.cassandra.CassandraContainer
+.. title:: testcontainers.cassandra.CassandraContainer

modules/cassandra/testcontainers/cassandra/__init__.py

@@ -0,0 +1,73 @@
+#
+#    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 os
+from typing import Optional
+
+from cassandra.cluster import Cluster, NoHostAvailable
+
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready
+
+
+class CassandraContainer(DockerContainer):
+    """
+    Cassandra database container.
+
+    Example:
+
+        .. doctest::
+
+            >>> from testcontainers.cassandra import CassandraContainer
+
+            >>> with CassandraContainer("cassandra:latest") as cassandra:
+            ...    db = cassandra.get_connection_client().test
+            ...    # Insert a database entry
+            ...    result = db.restaurants.insert_one(
+            ...        {
+            ...            "address": {
+            ...                "street": "2 Avenue",
+            ...                "zipcode": "10075",
+            ...                "building": "1480",
+            ...                "coord": [-73.9557413, 40.7720266]
+            ...            },
+            ...            "borough": "Manhattan",
+            ...            "cuisine": "Italian",
+            ...            "name": "Vella",
+            ...            "restaurant_id": "41704620"
+            ...        }
+            ...    )
+            ...    # Find the restaurant document
+            ...    cursor = db.restaurants.find({"borough": "Manhattan"})
+    """
+
+    def __init__(
+        self,
+        image: str = "cassandra:latest",
+        port: int = 9042,
+        **kwargs
+    ) -> None:
+        super().__init__(image=image, **kwargs)
+        self.port = port
+        self.with_exposed_ports(self.port)
+        self.cluster: Optional[Cluster] = None
+
+    @wait_container_is_ready(NoHostAvailable)
+    def _connect(self):
+        self.cluster = Cluster([self.get_container_host_ip()],
+                               port=self.get_exposed_port(self.port))
+        self.cluster.connect()
+
+    def start(self) -> "CassandraContainer":
+        super().start()
+        self._connect()
+        return self

modules/cassandra/tests/test_cassandra.py

@@ -0,0 +1,8 @@
+from testcontainers.cassandra import CassandraContainer
+
+
+def test_docker_run_cassandra():
+    with CassandraContainer("cassandra:4.1.4") as cassandra:
+        session = cassandra.cluster.connect()
+        result = session.execute("SELECT release_version FROM system.local;")
+        assert result.one().release_version == "4.1.4"

poetry.lock

@@ -29,6 +29,7 @@ packages = [
     { include = "testcontainers", from = "core" },
     { include = "testcontainers", from = "modules/arangodb" },
     { include = "testcontainers", from = "modules/azurite" },
+    { include = "testcontainers", from = "modules/cassandra" },
     { include = "testcontainers", from = "modules/clickhouse" },
     { include = "testcontainers", from = "modules/elasticsearch" },
     { include = "testcontainers", from = "modules/google" },
@@ -81,10 +82,12 @@ cx_Oracle = { version = "*", optional = true }
 pika = { version = "*", optional = true }
 redis = { version = "*", optional = true }
 selenium = { version = "*", optional = true }
+cassandra-driver = { version = "*", optional = true }
 
 [tool.poetry.extras]
 arangodb = ["python-arango"]
 azurite = ["azure-storage-blob"]
+cassandra = ["cassandra-driver"]
 clickhouse = ["clickhouse-driver"]
 elasticsearch = []
 google = ["google-cloud-pubsub"]
@@ -213,6 +216,7 @@ mypy_path = [
     "core",
 #    "modules/arangodb",
 #    "modules/azurite",
+#    "modules/cassandra",
 #    "modules/clickhouse",
 #    "modules/elasticsearch",
 #    "modules/google",