port cassandra from tc-java, fix small closing bug in both test

Author: alexanderankin

modules/cassandra/testcontainers/cassandra/__init__.py

@@ -1,6 +1,6 @@
 from testcontainers.core.config import MAX_TRIES
 from testcontainers.core.generic import DockerContainer
-from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+from testcontainers.core.waiting_utils import wait_for_logs
 
 
 class CassandraContainer(DockerContainer):
@@ -21,26 +21,30 @@ class CassandraContainer(DockerContainer):
         ...             "{'class': 'SimpleStrategy', 'replication_factor': '1'};")
     """
 
-    def __init__(self, image="rinscy/cassandra:latest", ports_to_expose=(9042,)):
+    def __init__(self, image="cassandra:latest", ports_to_expose=(9042,)):
         super().__init__(image)
         self.ports_to_expose = ports_to_expose
         self.with_exposed_ports(*self.ports_to_expose)
 
-    @wait_container_is_ready()
-    def _connect(self):
-        wait_for_logs(self, predicate="Starting listening for CQL clients", timeout=MAX_TRIES)
-        cluster = self.get_cluster()
-        cluster.connect()
+        self.with_env("CASSANDRA_SNITCH", "GossipingPropertyFileSnitch")
+        self.with_env("CASSANDRA_ENDPOINT_SNITCH", "GossipingPropertyFileSnitch")
+        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_DC", "datacenter1")
 
     def start(self):
         super().start()
-        self._connect()
+        wait_for_logs(self, predicate="Starting listening for CQL clients", timeout=MAX_TRIES)
         return self
 
     def get_cluster(self, **kwargs):
         from cassandra.cluster import Cluster
 
         container = self.get_wrapped_container()
         container.reload()
-        hostname = container.attrs["NetworkSettings"]["IPAddress"]
-        return Cluster(contact_points=[hostname], **kwargs)
+        hostname = self.get_container_host_ip()
+        port = self.get_exposed_port(9042)
+        return Cluster(contact_points=[hostname], port=port, **kwargs)

modules/cassandra/tests/test_cassandra.py

@@ -1,10 +1,12 @@
+import pytest
+
 from testcontainers.cassandra import CassandraContainer
 
 
-def test_docker_run_cassandra():
-    with CassandraContainer() as cassandra:
-        cluster = cassandra.get_cluster()
-        with cluster.connect() as session:
+@pytest.mark.parametrize("version", ["3.11.6", "4.1.4"], ids=["3.11.6", "4.1.4"])
+def test_docker_run_cassandra(version):
+    with CassandraContainer(image=f"cassandra:{version}") as cassandra:
+        with cassandra.get_cluster() as cluster, cluster.connect() as session:
             session.execute(
                 "CREATE KEYSPACE keyspace1 WITH replication = "
                 "{'class': 'SimpleStrategy', 'replication_factor': '1'};"

modules/scylla/tests/test_scylla.py

@@ -3,8 +3,7 @@
 
 def test_docker_run_scylla():
     with ScyllaContainer() as scylla:
-        cluster = scylla.get_cluster()
-        with cluster.connect() as session:
+        with scylla.get_cluster() as cluster, cluster.connect() as session:
             session.execute(
                 "CREATE KEYSPACE keyspace1 WITH replication = "
                 "{'class': 'SimpleStrategy', 'replication_factor': '1'};"