Improve test robustness

- Used shared test helpers more consistently

- Centralized port assignment to avoid conflicts

- Improve retry behaviour in get_connection

Author: lukakoning

conftest.py

@@ -0,0 +1,61 @@
+"""Repo-wide pytest configuration.
+
+This file applies to all tests in the repository (including package-level
+`*/tests` folders). It is intentionally minimal and focuses on integration-test
+hygiene.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+
+# Make shared fixtures/helpers from tests.integration_utils available everywhere.
+pytest_plugins = ["tests.integration_utils"]
+
+
+def _truthy_env(name: str) -> bool:
+    return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"}
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _cleanup_test_db_containers_session():
+    """Clean up ggmpilot test DB containers before and after the session.
+
+    CI failures due to "port is already allocated" or inability to start new
+    containers are often caused by leftover `*-docker-db-*` containers/volumes
+    from previous runs. This cleanup is best-effort and only runs when:
+
+    - Docker is reachable, and
+    - either `RUN_SLOW_TESTS` is enabled or we are in CI.
+
+    To opt out locally, set `GGMPILOT_KEEP_TEST_CONTAINERS=1`.
+    """
+
+    from tests.integration_utils import (
+        cleanup_all_test_db_containers,
+        docker_running,
+        slow_tests_enabled,
+    )
+
+    # In pytest-xdist, this session fixture would run in every worker process.
+    # Only run the global cleanup in the master/controller process to avoid
+    # workers deleting containers used by other workers.
+    if os.getenv("PYTEST_XDIST_WORKER"):
+        yield
+        return
+
+    if _truthy_env("GGMPILOT_KEEP_TEST_CONTAINERS"):
+        yield
+        return
+
+    should_cleanup = docker_running() and (slow_tests_enabled() or _truthy_env("CI"))
+    if should_cleanup:
+        cleanup_all_test_db_containers()
+
+    yield
+
+    if should_cleanup:
+        cleanup_all_test_db_containers()

dev_sql_server/get_connection.py

@@ -29,7 +29,13 @@
 # Driver‑specific helpers
 # ────────────────────────────────────────────────────────────────────────────────
 def _connect_postgres(cfg: Dict[str, Any]):
-    return psycopg2.connect(**cfg)
+    # psycopg2 can block for a while on connect() if no connect_timeout is set.
+    # In CI (and especially under parallel starts) we want fast retries.
+    cfg2 = dict(cfg)
+    cfg2.setdefault(
+        "connect_timeout", int(os.getenv("GGMPILOT_DB_CONNECT_TIMEOUT", "5"))
+    )
+    return psycopg2.connect(**cfg2)
 
 
 def _connect_oracle(cfg: Dict[str, Any]):
@@ -44,11 +50,17 @@ def _connect_oracle(cfg: Dict[str, Any]):
 
 
 def _connect_mssql(cfg: Dict[str, Any]):
+    # Keep individual connection attempts short so _wait_for_db_ready can retry.
+    # ODBC driver 18 defaults to Encrypt=yes; TrustServerCertificate is required
+    # for local/dev containers with self-signed certs.
+    connect_timeout = int(os.getenv("GGMPILOT_DB_CONNECT_TIMEOUT", "5"))
     conn_str = (
         f"DRIVER={{{SQL_SERVER_DRIVER}}};"
         f"SERVER={cfg['host']},{cfg['port']};"
         f"DATABASE={cfg['dbname']};"
         f"UID={cfg['user']};PWD={cfg['password']};"
+        f"Connection Timeout={connect_timeout};"
+        "Encrypt=yes;"
         "TrustServerCertificate=yes;"
     )
     return pyodbc.connect(conn_str)
@@ -500,11 +512,19 @@ def get_connection(
     container_name = container_name or f"{db_type}-docker-db-{port_effective}"
     volume_name = volume_name or f"{container_name}_data"
 
+    user_supplied_max_wait = max_wait_seconds is not None
+
     if max_wait_seconds is None:
+        # Defaults are intentionally conservative for CI stability.
+        # Callers can always override via max_wait_seconds.
         if db_type == "oracle":
             max_wait_seconds = 600
         elif db_type in ("mysql", "mariadb"):
             max_wait_seconds = 180
+        elif db_type == "postgres":
+            max_wait_seconds = 180
+        elif db_type == "mssql":
+            max_wait_seconds = 360
         else:
             max_wait_seconds = 120
 
@@ -527,6 +547,24 @@ def get_connection(
         container_force_refresh,
     )
 
+    # If the container was newly created, DB initialization can take
+    # substantially longer on cold CI runners (esp. MSSQL). Only adjust
+    # when the caller didn't explicitly choose a timeout.
+    if (not user_supplied_max_wait) and was_created:
+        if db_type == "postgres":
+            max_wait_seconds = max(max_wait_seconds, 240)
+        elif db_type == "mssql":
+            max_wait_seconds = max(max_wait_seconds, 480)
+
+    # CI runners can be slower (image pulls, constrained IO). Allow opt-in
+    # override while keeping local runs reasonable.
+    if not user_supplied_max_wait:
+        if os.getenv("CI", "").strip().lower() in {"1", "true", "yes", "on"}:
+            if db_type == "postgres":
+                max_wait_seconds = max(max_wait_seconds, 240)
+            elif db_type == "mssql":
+                max_wait_seconds = max(max_wait_seconds, 480)
+
     # Prepare configs for master and target DB
     # When running inside Docker, localhost refers to the container itself.
     # Prefer host.docker.internal; if not resolvable on Linux, fall back to Docker bridge gateway.
@@ -537,7 +575,9 @@ def get_connection(
         except Exception:
             host_addr = os.getenv("HOST_GATEWAY_IP", "172.17.0.1")
     else:
-        host_addr = "localhost"
+        # Prefer IPv4 loopback to avoid occasional localhost/IPv6 resolution
+        # differences across CI runners.
+        host_addr = "127.0.0.1"
 
     # Choose appropriate admin DB/user per backend
     if db_type == "postgres":

odata_to_staging/tests/test_integration_odata_to_staging_postgres.py

@@ -3,32 +3,18 @@
 import runpy
 import io
 import contextlib
-import subprocess
 
 import pytest
 import requests
 from sqlalchemy import text
 
 from dev_sql_server.get_connection import get_connection
+from tests.integration_utils import docker_running, slow_tests_enabled, ports_dest
 
 
 NORTHWIND_V2 = "https://services.odata.org/V2/Northwind/Northwind.svc/"
 
 
-def _docker_running() -> bool:
-    try:
-        res = subprocess.run(
-            ["docker", "info"], capture_output=True, text=True, timeout=5
-        )
-        return res.returncode == 0
-    except Exception:
-        return False
-
-
-def _slow_tests_enabled() -> bool:
-    return os.getenv("RUN_SLOW_TESTS", "0").lower() in {"1", "true", "yes", "on"}
-
-
 def _northwind_available() -> bool:
     """Best-effort sanity check that the Northwind service is reachable and returns JSON for an entity set.
 
@@ -57,25 +43,27 @@ def _northwind_available() -> bool:
 @pytest.mark.slow
 @pytest.mark.postgres
 @pytest.mark.skipif(
-    not _slow_tests_enabled(),
+    not slow_tests_enabled(),
     reason="RUN_SLOW_TESTS not enabled; set to 1 to run slow integration tests.",
 )
 @pytest.mark.skipif(
-    not _docker_running(),
+    not docker_running(),
     reason="Docker is not available/running; required for this integration test.",
 )
 @pytest.mark.skipif(
     not _northwind_available(),
     reason="Northwind OData service not reachable; skipping to avoid flakey failures.",
 )
 def test_main_odata_to_staging_postgres(tmp_path):
+    dst_port = ports_dest["postgres"]
+
     # Start a fresh Postgres destination
     engine = get_connection(
         db_type="postgres",
         db_name="ggm_odata_to_staging",
         user="sa",
         password="S3cureP@ssw0rd!23243",
-        port=5434,
+        port=dst_port,
         force_refresh=True,
         print_tables=False,
     )
@@ -95,7 +83,7 @@ def test_main_odata_to_staging_postgres(tmp_path):
 DST_USERNAME=sa
 DST_PASSWORD=S3cureP@ssw0rd!23243
 DST_HOST=localhost
-DST_PORT=5434
+DST_PORT={dst_port}
 DST_DB=ggm_odata_to_staging
 DST_SCHEMA=staging
 

sql_to_staging/tests/test_integration_direct_transfer_streaming.py

@@ -4,74 +4,37 @@
 
 import logging
 import os
-import shutil
-import subprocess
 import re
 
 import pytest
-import docker
 from sqlalchemy import text
 from dotenv import load_dotenv
 
 from dev_sql_server.get_connection import get_connection
 from sql_to_staging.functions.direct_transfer import direct_transfer
+from tests.integration_utils import (
+    cleanup_db_container_by_port,
+    docker_running,
+    ports,
+    ports_dest,
+    slow_tests_enabled,
+)
 
 
 load_dotenv("tests/.env")
 
 
-def _docker_running() -> bool:
-    if not shutil.which("docker"):
-        return False
-    try:
-        res = subprocess.run(
-            ["docker", "info"], capture_output=True, text=True, timeout=5
-        )
-        return res.returncode == 0
-    except Exception:
-        return False
-
-
-def _slow_tests_enabled() -> bool:
-    return os.getenv("RUN_SLOW_TESTS", "0").lower() in {"1", "true", "yes", "on"}
-
-
 # Use explicit ports consistent with other integration tests
-SRC_PORT = 5433
-DST_PORT = 5434
-
-
-def _cleanup_db_containers(name: str, port: int):
-    """Stop and remove the container/volume our get_connection uses for a db/port pair."""
-    client = docker.from_env()
-    cname = f"{name}-docker-db-{port}"
-    try:
-        c = client.containers.get(cname)
-        try:
-            c.stop()
-        except Exception:
-            pass
-        try:
-            c.remove()
-        except Exception:
-            pass
-    except Exception:
-        pass
-    # remove associated volume
-    vname = f"{cname}_data"
-    try:
-        v = client.volumes.get(vname)
-        v.remove(force=True)
-    except Exception:
-        pass
+SRC_PORT = ports["postgres"]
+DST_PORT = ports_dest["postgres"]
 
 
 @pytest.mark.skipif(
-    not _slow_tests_enabled(),
+    not slow_tests_enabled(),
     reason="RUN_SLOW_TESTS not enabled; set to 1 to run slow integration tests.",
 )
 @pytest.mark.skipif(
-    not _docker_running(),
+    not docker_running(),
     reason="Docker is not available/running; required for this integration test.",
 )
 @pytest.mark.parametrize(
@@ -93,8 +56,8 @@ def test_direct_transfer_streams_in_chunks_postgres(
     table = "stream_check"
 
     # Ensure a clean slate before starting
-    _cleanup_db_containers("postgres", SRC_PORT)
-    _cleanup_db_containers("postgres", DST_PORT)
+    cleanup_db_container_by_port("postgres", SRC_PORT)
+    cleanup_db_container_by_port("postgres", DST_PORT)
 
     try:
         # Start Postgres source and create table with many rows
@@ -164,5 +127,5 @@ def test_direct_transfer_streams_in_chunks_postgres(
                     inserts.append((int(m.group(1)), int(m.group(2))))
         assert inserts == expected_batches
     finally:
-        _cleanup_db_containers("postgres", SRC_PORT)
-        _cleanup_db_containers("postgres", DST_PORT)
+        cleanup_db_container_by_port("postgres", SRC_PORT)
+        cleanup_db_container_by_port("postgres", DST_PORT)

sql_to_staging/tests/test_upload_parquet_mssql_datetime2.py

@@ -13,47 +13,26 @@
 
 from dev_sql_server.get_connection import get_connection
 from utils.parquet.upload_parquet import upload_parquet
-
-
-def _docker_running() -> bool:
-    import subprocess
-
-    try:
-        res = subprocess.run(
-            ["docker", "info"], capture_output=True, text=True, timeout=5
-        )
-        return res.returncode == 0
-    except Exception:
-        return False
-
-
-def _slow_tests_enabled() -> bool:
-    import os
-
-    return os.getenv("RUN_SLOW_TESTS", "0").lower() in {"1", "true", "yes", "on"}
-
-
-def _mssql_driver_available() -> bool:
-    try:
-        import pyodbc  # noqa: F401
-
-        return any("ODBC Driver 18 for SQL Server" in d for d in pyodbc.drivers())
-    except Exception:
-        return False
+from tests.integration_utils import (
+    docker_running,
+    mssql_driver_available,
+    port_for_worker,
+    slow_tests_enabled,
+)
 
 
 @pytest.mark.slow
 @pytest.mark.mssql
 @pytest.mark.skipif(
-    not _slow_tests_enabled(),
+    not slow_tests_enabled(),
     reason="RUN_SLOW_TESTS not enabled; set to 1 to run slow integration tests.",
 )
 @pytest.mark.skipif(
-    not _docker_running(),
+    not docker_running(),
     reason="Docker is not available/running; required for this integration test.",
 )
 @pytest.mark.skipif(
-    not _mssql_driver_available(),
+    not mssql_driver_available(),
     reason="ODBC Driver 18 for SQL Server not installed; required for MSSQL test.",
 )
 def test_upload_parquet_mssql_uses_datetime2(tmp_path: Path):
@@ -63,7 +42,9 @@ def test_upload_parquet_mssql_uses_datetime2(tmp_path: Path):
         db_name="ggm_upload_parquet_dt2",
         user="sa",
         password="S3cureP@ssw0rd!23243",
-        port=1436,  # use a distinct port to avoid clashes with other tests
+        port=port_for_worker(
+            1436
+        ),  # use a distinct port to avoid clashes with other tests
         force_refresh=True,
         sql_folder=None,
         sql_suffix_filter=True,