adjustments

modules/cockroachdb/testcontainers/__init__.py → modules/cockroachdb/testcontainers/cockroachdb/__init__.py

@@ -12,8 +12,11 @@
 #    under the License.
 from os import environ
 from typing import Optional
+from urllib.error import HTTPError, URLError
+from urllib.request import urlopen
 
 from testcontainers.core.generic import DbContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
 
 
 class CockroachDBContainer(DbContainer):
@@ -32,28 +35,29 @@ class CockroachDBContainer(DbContainer):
             >>> import sqlalchemy
             >>> from testcontainers.cockroachdb import CockroachDBContainer
 
-            >>> with CockroachDBContainer('cockroachdb/cockroach:latest') as crdb:
+            >>> with CockroachDBContainer('cockroachdb/cockroach:v24.1.1') as crdb:
             ...     engine = sqlalchemy.create_engine(crdb.get_connection_url())
             ...     with engine.begin() as connection:
             ...         result = connection.execute(sqlalchemy.text("select version()"))
             ...         version, = result.fetchone()
 
     """
 
+    COCKROACH_DB_PORT: int = 26257
+    COCKROACH_API_PORT: int = 8080
+
     def __init__(
         self,
-        image: str = "cockroachdb/cockroach:latest",
+        image: str = "cockroachdb/cockroach:v24.1.1",
         username: Optional[str] = None,
         password: Optional[str] = None,
         dbname: Optional[str] = None,
-        port: int = 26257,
         dialect="cockroachdb+psycopg2",
         **kwargs,
     ) -> None:
         super().__init__(image, **kwargs)
 
-        self.port = port
-        self.with_exposed_ports(self.port)
+        self.with_exposed_ports(self.COCKROACH_DB_PORT, self.COCKROACH_API_PORT)
         self.username = username or environ.get("COCKROACH_USER", "cockroach")
         self.password = password or environ.get("COCKROACH_PASSWORD", "arthropod")
         self.dbname = dbname or environ.get("COCKROACH_DATABASE", "roach")
@@ -69,9 +73,25 @@ def _configure(self) -> None:
             cmd += " --insecure"
         self.with_command(cmd)
 
+    @wait_container_is_ready(HTTPError, URLError)
+    def _connect(self) -> None:
+        host = self.get_container_host_ip()
+        url = f"http://{host}:{self.get_exposed_port(self.COCKROACH_API_PORT)}/health"
+        self._wait_for_health(url)
+        wait_for_logs(self, "finished creating default user*")
+
+    @staticmethod
+    def _wait_for_health(url):
+        with urlopen(url) as response:
+            response.read()
+
     def get_connection_url(self) -> str:
         conn_str = super()._create_connection_url(
-            dialect=self.dialect, username=self.username, password=self.password, dbname=self.dbname, port=self.port
+            dialect=self.dialect,
+            username=self.username,
+            password=self.password,
+            dbname=self.dbname,
+            port=self.COCKROACH_DB_PORT,
         )
 
         if self.password:

modules/cockroachdb/tests/test_cockroachdb.py

@@ -1,18 +1,14 @@
-from pathlib import Path
-import re
-from unittest import mock
-
-import pytest
 import sqlalchemy
 
-from testcontainers.core.utils import is_arm
 from testcontainers.cockroachdb import CockroachDBContainer
 
+
 def test_docker_run_mysql():
-    config = CockroachDBContainer("cockroachdb/cockroach:24.0.1")
+    config = CockroachDBContainer("cockroachdb/cockroach:v24.1.1")
     with config as crdb:
         engine = sqlalchemy.create_engine(crdb.get_connection_url())
         with engine.begin() as connection:
             result = connection.execute(sqlalchemy.text("select version()"))
             for row in result:
-                assert row[0].startswith("24.0.1")
+                assert "CockroachDB" in row[0]
+                assert "v24.1.1" in row[0]

pyproject.toml

@@ -108,6 +108,7 @@ arangodb = ["python-arango"]
 azurite = ["azure-storage-blob"]
 cassandra = []
 clickhouse = ["clickhouse-driver"]
+cockroachdb = []
 elasticsearch = []
 google = ["google-cloud-pubsub", "google-cloud-datastore"]
 influxdb = ["influxdb", "influxdb-client"]
@@ -158,6 +159,7 @@ hvac = "2.1.0"
 pymilvus = "2.4.3"
 httpx = "0.27.0"
 paho-mqtt = "2.1.0"
+sqlalchemy-cockroachdb = "2.0.2"
 
 [[tool.poetry.source]]
 name = "PyPI"