add support async postgres driver

core/testcontainers/core/generic.py

@@ -28,13 +28,16 @@ class DbContainer(DockerContainer):
     """
     Generic database container.
     """
+
+    DEFAULT_DRIVER = None
+
     @wait_container_is_ready(*ADDITIONAL_TRANSIENT_ERRORS)
-    def _connect(self) -> None:
+    def _connect(self, driver: Optional[str] = None) -> None:
         import sqlalchemy
-        engine = sqlalchemy.create_engine(self.get_connection_url())
+        engine = sqlalchemy.create_engine(self.get_connection_url(driver=driver))
         engine.connect()
 
-    def get_connection_url(self) -> str:
+    def get_connection_url(self, host=None, driver=None) -> str:
         raise NotImplementedError
 
     def _create_connection_url(self, dialect: str, username: str, password: str,
@@ -54,7 +57,7 @@ def _create_connection_url(self, dialect: str, username: str, password: str,
     def start(self) -> 'DbContainer':
         self._configure()
         super().start()
-        self._connect()
+        self._connect(self.DEFAULT_DRIVER)
         return self
 
     def _configure(self) -> None:

postgres/setup.py

@@ -14,6 +14,8 @@
         "testcontainers-core",
         "sqlalchemy",
         "psycopg2-binary",
+        "asyncpg",
+        "pytest-asyncio",
     ],
     python_requires=">=3.7",
 )

postgres/testcontainers/postgres/__init__.py

@@ -41,6 +41,7 @@ class PostgresContainer(DbContainer):
     POSTGRES_USER = os.environ.get("POSTGRES_USER", "test")
     POSTGRES_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "test")
     POSTGRES_DB = os.environ.get("POSTGRES_DB", "test")
+    DEFAULT_DRIVER = "psycopg2"
 
     def __init__(self, image: str = "postgres:latest", port: int = 5432, user: Optional[str] = None,
                  password: Optional[str] = None, dbname: Optional[str] = None,
@@ -59,9 +60,12 @@ def _configure(self) -> None:
         self.with_env("POSTGRES_PASSWORD", self.POSTGRES_PASSWORD)
         self.with_env("POSTGRES_DB", self.POSTGRES_DB)
 
-    def get_connection_url(self, host=None) -> str:
+    def get_connection_url(self, host=None, driver: str = None) -> str:
+        if not driver:
+            driver = self.driver
+
         return super()._create_connection_url(
-            dialect="postgresql+{}".format(self.driver), username=self.POSTGRES_USER,
+            dialect="postgresql+{}".format(driver), username=self.POSTGRES_USER,
             password=self.POSTGRES_PASSWORD, db_name=self.POSTGRES_DB, host=host,
             port=self.port_to_expose,
         )

postgres/tests/test_postgres.py

@@ -1,4 +1,6 @@
+import pytest
 import sqlalchemy
+from sqlalchemy.ext.asyncio import create_async_engine
 from testcontainers.postgres import PostgresContainer
 
 
@@ -18,3 +20,13 @@ def test_docker_run_postgres_with_driver_pg8000():
         engine = sqlalchemy.create_engine(postgres.get_connection_url())
         with engine.begin() as connection:
             connection.execute(sqlalchemy.text("select 1=1"))
+
+
+@pytest.mark.asyncio
+async def test_docker_run_async_postgres():
+    with PostgresContainer("postgres:9.5", driver="asyncpg") as postgres:
+        engine = create_async_engine(postgres.get_connection_url())
+        async with engine.begin() as connection:
+            result = await connection.execute(sqlalchemy.text("select version()"))
+            for row in result:
+                assert row[0].lower().startswith("postgresql 9.5")