From cfa7e856eb6fcb6d30019de63fa67d1fd472613f Mon Sep 17 00:00:00 2001
From: Oliver Lambson <oliverlambson@gmail.com>
Date: Fri, 28 Jun 2024 18:54:33 +0100
Subject: [PATCH] fix(postgres): get_connection_url(driver=None) should return
 postgres://... (#588)

Fixes #587
---
 .../testcontainers/postgres/__init__.py       |  2 +-
 modules/postgres/tests/test_postgres.py       | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/modules/postgres/testcontainers/postgres/__init__.py b/modules/postgres/testcontainers/postgres/__init__.py
index 9b347aa6..80baef75 100644
--- a/modules/postgres/testcontainers/postgres/__init__.py
+++ b/modules/postgres/testcontainers/postgres/__init__.py
@@ -79,7 +79,7 @@ def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] =
         driver. The optional driver argument to :code:`get_connection_url` overwrites the constructor
         set value. Pass :code:`driver=None` to get URLs without a driver.
         """
-        driver_str = self.driver if driver is _UNSET else f"+{driver}"
+        driver_str = "" if driver is None else self.driver if driver is _UNSET else f"+{driver}"
         return super()._create_connection_url(
             dialect=f"postgresql{driver_str}",
             username=self.username,
diff --git a/modules/postgres/tests/test_postgres.py b/modules/postgres/tests/test_postgres.py
index 52840361..38c856bf 100644
--- a/modules/postgres/tests/test_postgres.py
+++ b/modules/postgres/tests/test_postgres.py
@@ -97,3 +97,27 @@ def test_show_how_to_initialize_db_via_initdb_dir():
             result = result.fetchall()
             assert len(result) == 1
             assert result[0] == (1, "sally", "sells seashells")
+
+
+def test_none_driver_urls():
+    user = "root"
+    password = "pass"
+    kwargs = {
+        "username": user,
+        "password": password,
+    }
+    with PostgresContainer("postgres:16-alpine", driver=None, **kwargs) as container:
+        port = container.get_exposed_port(5432)
+        host = container.get_container_host_ip()
+        expected_url = f"postgresql://{user}:{password}@{host}:{port}/test"
+
+        url = container.get_connection_url()
+        assert url == expected_url
+
+    with PostgresContainer("postgres:16-alpine", **kwargs) as container:
+        port = container.get_exposed_port(5432)
+        host = container.get_container_host_ip()
+        expected_url = f"postgresql://{user}:{password}@{host}:{port}/test"
+
+        url = container.get_connection_url(driver=None)
+        assert url == expected_url