Add ResourceCleaner for stop containers after exit (#104)

Author: Mark Andreev

testcontainers/core/cleaner.py

@@ -0,0 +1,36 @@
+import atexit
+
+from docker.errors import NotFound
+
+from testcontainers.core.utils import setup_logger
+
+logger = setup_logger(__name__)
+
+
+class ResourceCleaner:
+    __INSTANCE = None
+
+    def __init__(self):
+        self.containers = []
+
+    def attach(self, container):
+        self.containers.append(container)
+
+        return container
+
+    def clean(self):
+        for c in self.containers:
+            try:
+                c.stop()
+            except NotFound:
+                pass
+            except Exception as e:
+                logger.exception(e)
+
+    @classmethod
+    def instance(cls):
+        if cls.__INSTANCE is None:
+            cls.__INSTANCE = ResourceCleaner()
+            atexit.register(cls.__INSTANCE.clean)
+
+        return cls.__INSTANCE

testcontainers/core/docker_client.py

@@ -14,6 +14,8 @@
 import urllib
 import docker
 from docker.models.containers import Container
+
+from testcontainers.core.cleaner import ResourceCleaner
 from testcontainers.core.utils import inside_container
 from testcontainers.core.utils import default_gateway_ip
 
@@ -30,15 +32,17 @@ def run(self, image: str,
             stdout: bool = True,
             stderr: bool = False,
             remove: bool = False, **kwargs) -> Container:
-        return self.client.containers.run(image,
-                                          command=command,
-                                          stdout=stdout,
-                                          stderr=stderr,
-                                          remove=remove,
-                                          detach=detach,
-                                          environment=environment,
-                                          ports=ports,
-                                          **kwargs)
+        return ResourceCleaner.instance().attach(
+            self.client.containers.run(image,
+                                       command=command,
+                                       stdout=stdout,
+                                       stderr=stderr,
+                                       remove=remove,
+                                       detach=detach,
+                                       environment=environment,
+                                       ports=ports,
+                                       **kwargs)
+        )
 
     def port(self, container_id, port):
         port_mappings = self.client.api.port(container_id, port)