Testing: Generalize KeepaliveContainer from CrateDBContainer

cratedb_toolkit/testing/testcontainers/cratedb.py

@@ -18,10 +18,12 @@
 from testcontainers.core.generic import DbContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
+from cratedb_toolkit.testing.testcontainers.util import KeepaliveContainer, asbool
+
 logger = logging.getLogger(__name__)
 
 
-class CrateDBContainer(DbContainer):
+class CrateDBContainer(KeepaliveContainer, DbContainer):
     """
     CrateDB database container.
 
@@ -48,21 +50,20 @@ class CrateDBContainer(DbContainer):
     CRATEDB_USER = os.environ.get("CRATEDB_USER", "crate")
     CRATEDB_PASSWORD = os.environ.get("CRATEDB_PASSWORD", "")
     CRATEDB_DB = os.environ.get("CRATEDB_DB", "doc")
-    CRATEDB_KEEPALIVE = os.environ.get("CRATEDB_KEEPALIVE", os.environ.get("TC_KEEPALIVE", False))
+    KEEPALIVE = asbool(os.environ.get("CRATEDB_KEEPALIVE", os.environ.get("TC_KEEPALIVE", False)))
 
     # TODO: Dual-port use with 4200+5432.
     def __init__(
         self,
-        image: str = "crate:latest",
+        image: str = "crate/crate:nightly",
         port: int = 4200,
         user: Optional[str] = None,
         password: Optional[str] = None,
         dbname: Optional[str] = None,
         dialect: str = "crate",
-        keepalive: bool = False,
         **kwargs,
     ) -> None:
-        super(CrateDBContainer, self).__init__(image=image, **kwargs)
+        super().__init__(image=image, **kwargs)
 
         self._name = "testcontainers-cratedb"  # -{os.getpid()}
         self._command = "-Cdiscovery.type=single-node -Ccluster.routing.allocation.disk.threshold_enabled=false"
@@ -74,13 +75,11 @@ def __init__(
         self.CRATEDB_PASSWORD = password or self.CRATEDB_PASSWORD
         self.CRATEDB_DB = dbname or self.CRATEDB_DB
 
-        self.keepalive = keepalive or self.CRATEDB_KEEPALIVE
         self.port_to_expose = port
         self.dialect = dialect
 
-        self.with_exposed_ports(self.port_to_expose)
-
     def _configure(self) -> None:
+        self.with_exposed_ports(self.port_to_expose)
         self.with_env("CRATEDB_USER", self.CRATEDB_USER)
         self.with_env("CRATEDB_PASSWORD", self.CRATEDB_PASSWORD)
         self.with_env("CRATEDB_DB", self.CRATEDB_DB)
@@ -100,46 +99,3 @@ def _connect(self):
         # TODO: Better use a network connectivity health check?
         #       In `testcontainers-java`, there is the `HttpWaitStrategy`.
         wait_for_logs(self, predicate="o.e.n.Node.*started", timeout=MAX_TRIES)
-
-    def start(self):
-        """
-        Improved `start()` method, supporting service-keepalive.
-
-        In order to keep the service running where it normally would be torn down,
-        define the `CRATEDB_KEEPALIVE` or `TC_KEEPALIVE` environment variables.
-        """
-
-        self._configure()
-
-        logger.info("Pulling image %s", self.image)
-        docker_client = self.get_docker_client()
-
-        # Check if container is already running, and whether it should be reused.
-        containers_running = docker_client.client.api.containers(all=True, filters={"name": self._name})
-        start_container = not containers_running
-
-        if start_container:
-            logger.info("Starting CrateDB")
-            self._container = docker_client.run(
-                self.image,
-                command=self._command,
-                detach=True,
-                environment=self.env,
-                ports=self.ports,
-                name=self._name,
-                volumes=self.volumes,
-                **self._kwargs,
-            )
-        else:
-            container_id = containers_running[0]["Id"]
-            self._container = docker_client.client.containers.get(container_id)
-
-        logger.info("Container started: %s", self._container.short_id)
-        self._connect()
-        return self
-
-    def stop(self, **kwargs):
-        if not self.keepalive:
-            logger.info("Stopping CrateDB")
-            return super().stop()
-        return None

cratedb_toolkit/testing/testcontainers/util.py

@@ -10,8 +10,28 @@
 #    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 logging
+import os
+from typing import Any
+
 from testcontainers.core.container import DockerContainer
 
+logger = logging.getLogger(__name__)
+
+
+# from sqlalchemy.util.langhelpers
+# from paste.deploy.converters
+def asbool(obj: Any) -> bool:
+    if isinstance(obj, str):
+        obj = obj.strip().lower()
+        if obj in ["true", "yes", "on", "y", "t", "1"]:
+            return True
+        elif obj in ["false", "no", "off", "n", "f", "0"]:
+            return False
+        else:
+            raise ValueError("String is not true/false: %r" % obj)
+    return bool(obj)
+
 
 class ExtendedDockerContainer(DockerContainer):
     """
@@ -34,3 +54,81 @@ def get_real_host_address(self) -> str:
         For example, `172.17.0.4:9000`.
         """
         return f"{self.get_real_host_ip()}:{self.port_to_expose}"
+
+
+class KeepaliveContainer(DockerContainer):
+    """
+    Improved `start()`/`stop()` methods, supporting service-keepalive.
+
+    In order to keep the service running where it normally would be torn down,
+    define the `TC_KEEPALIVE` environment variable.
+    """
+
+    KEEPALIVE = asbool(os.environ.get("TC_KEEPALIVE", False))
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ) -> None:
+        self.keepalive = self.KEEPALIVE
+        if "keepalive" in kwargs:
+            self.keepalive = kwargs["keepalive"]
+            del kwargs["keepalive"]
+        super().__init__(*args, **kwargs)
+
+    def start(self):
+        """
+        Improved `start()` method, supporting service-keepalive.
+
+        In order to keep the service running where it normally would be torn down,
+        define the `CRATEDB_KEEPALIVE` or `TC_KEEPALIVE` environment variables.
+        """
+
+        self._configure()
+
+        if self._name is None:
+            raise ValueError(
+                "KeepaliveContainer does not support unnamed containers. Use `.with_name()` to assign a name."
+            )
+
+        docker_client = self.get_docker_client()
+
+        # Check if container is already running, and whether it should be reused.
+        logger.info(f"Searching for container: {self._name}")
+        containers = docker_client.client.api.containers(all=True, filters={"name": self._name})
+
+        if not containers:
+            logger.info(f"Creating container from image: {self.image}")
+            self._container = docker_client.run(
+                self.image,
+                command=self._command,
+                detach=True,
+                environment=self.env,
+                ports=self.ports,
+                name=self._name,
+                volumes=self.volumes,
+                **self._kwargs,
+            )
+            logger.info(f"Container created: {self._container.name}")
+        else:
+            container_id = containers[0]["Id"]
+            container_names = containers[0]["Names"]
+            logger.info(f"Found container for reuse: {container_id} ({container_names})")
+            self._container = docker_client.client.containers.get(container_id)
+            container_name = self._container.name
+            if self._container.status != "running":
+                logger.info(f"Starting container: {container_id} ({container_name})")
+                self._container.start()
+
+        self._connect()
+        return self
+
+    def stop(self, **kwargs):
+        """
+        Shut down container again, unless "keepalive" is enabled.
+        """
+        if not self.keepalive:
+            logger.info("Stopping container")
+            return super().stop()
+        return None