fix: Add CassandraContainer (#476)

Co-authored-by: David Ankin <[email protected]>

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/chroma/README
     modules/clickhouse/README
     modules/elasticsearch/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,62 @@
+#
+#    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 testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
+
+
+class CassandraContainer(DockerContainer):
+    """
+    Cassandra database container.
+
+    Example:
+
+        .. doctest::
+
+            >>> from testcontainers.cassandra import CassandraContainer
+            >>> from cassandra.cluster import Cluster, DCAwareRoundRobinPolicy
+
+            >>> with CassandraContainer("cassandra:4.1.4") as cassandra, Cluster(
+            ...    cassandra.get_contact_points(),
+            ...    load_balancing_policy=DCAwareRoundRobinPolicy(cassandra.get_local_datacenter()),
+            ... ) as cluster:
+            ...    session = cluster.connect()
+            ...    result = session.execute("SELECT release_version FROM system.local;")
+            ...    result.one().release_version
+            '4.1.4'
+    """
+
+    CQL_PORT = 9042
+    DEFAULT_LOCAL_DATACENTER = "datacenter1"
+
+    def __init__(self, image: str = "cassandra:latest", **kwargs) -> None:
+        super().__init__(image=image, **kwargs)
+        self.with_exposed_ports(self.CQL_PORT)
+        self.with_env("JVM_OPTS", "-Dcassandra.skip_wait_for_gossip_to_settle=0 -Dcassandra.initial_token=0")
+        self.with_env("HEAP_NEWSIZE", "128M")
+        self.with_env("MAX_HEAP_SIZE", "1024M")
+        self.with_env("CASSANDRA_ENDPOINT_SNITCH", "GossipingPropertyFileSnitch")
+        self.with_env("CASSANDRA_DC", self.DEFAULT_LOCAL_DATACENTER)
+
+    def _connect(self):
+        wait_for_logs(self, "Startup complete")
+
+    def start(self) -> "CassandraContainer":
+        super().start()
+        self._connect()
+        return self
+
+    def get_contact_points(self) -> list[tuple[str, int]]:
+        return [(self.get_container_host_ip(), int(self.get_exposed_port(self.CQL_PORT)))]
+
+    def get_local_datacenter(self) -> str:
+        return self.env.get("CASSANDRA_DC", self.DEFAULT_LOCAL_DATACENTER)

modules/cassandra/tests/test_cassandra.py

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

poetry.lock

@@ -30,6 +30,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/chroma" },
     { include = "testcontainers", from = "modules/clickhouse" },
     { include = "testcontainers", from = "modules/elasticsearch" },
@@ -93,6 +94,7 @@ chromadb-client = { version = "*", optional = true }
 [tool.poetry.extras]
 arangodb = ["python-arango"]
 azurite = ["azure-storage-blob"]
+cassandra = ["cassandra-driver"]
 clickhouse = ["clickhouse-driver"]
 elasticsearch = []
 google = ["google-cloud-pubsub", "google-cloud-datastore"]
@@ -130,6 +132,7 @@ pg8000 = "*"
 sqlalchemy = "*"
 psycopg = "*"
 kafka-python = "^2.0.2"
+cassandra-driver = "*"
 
 [[tool.poetry.source]]
 name = "PyPI"
@@ -228,6 +231,7 @@ mypy_path = [
     "core",
 #    "modules/arangodb",
 #    "modules/azurite",
+#    "modules/cassandra",
 #    "modules/clickhouse",
 #    "modules/elasticsearch",
 #    "modules/google",