migrate alloydb

Author: provinzkraut

src/pytest_databases/docker/alloydb_omni.py

@@ -1,136 +1,44 @@
 from __future__ import annotations
 
-import os
-import sys
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 import psycopg
 import pytest
 
+from pytest_databases._service import DockerService
 from pytest_databases.docker import DockerServiceRegistry
-from pytest_databases.helpers import simple_string_hash
-
-if TYPE_CHECKING:
-    from collections.abc import Generator
-
-
-COMPOSE_PROJECT_NAME: str = f"pytest-databases-alloydb-{simple_string_hash(__file__)}"
-
-
-def alloydb_omni_responsive(host: str, port: int, user: str, password: str, database: str) -> bool:
-    try:
-        conn = psycopg.connect(
-            host=host,
-            port=port,
-            user=user,
-            database=database,
-            password=password,
-        )
-    except Exception:  # noqa: BLE001
-        return False
-
-    try:
-        db_open = conn.fetchrow("SELECT 1")
-        return bool(db_open is not None and db_open[0] == 1)
-    finally:
-        conn.close()
-
-
-@pytest.fixture(scope="session")
-def alloydb_compose_project_name() -> str:
-    return os.environ.get("COMPOSE_PROJECT_NAME", COMPOSE_PROJECT_NAME)
-
-
-@pytest.fixture(autouse=False, scope="session")
-def alloydb_docker_services(
-    alloydb_compose_project_name: str, worker_id: str = "main"
-) -> Generator[DockerServiceRegistry, None, None]:
-    if os.getenv("GITHUB_ACTIONS") == "true" and sys.platform != "linux":
-        pytest.skip("Docker not available on this platform")
-
-    with DockerServiceRegistry(worker_id, compose_project_name=alloydb_compose_project_name) as registry:
-        yield registry
+from pytest_databases.docker.postgres import (
+    _make_connection_string,
+    _provide_postgres_service,
+    PostgresService as AlloyDBService,
+)
 
+__all__ = ("AlloyDBService",)
 
-@pytest.fixture(scope="session")
-def postgres_user() -> str:
-    return "postgres"
 
-
-@pytest.fixture(scope="session")
-def postgres_password() -> str:
-    return "super-secret"
-
-
-@pytest.fixture(scope="session")
-def postgres_database() -> str:
-    return "postgres"
-
-
-@pytest.fixture(scope="session")
-def alloydb_omni_port() -> int:
-    return 5420
-
-
-@pytest.fixture(scope="session")
-def default_alloydb_omni_service_name() -> str:
-    return "alloydb"
-
-
-@pytest.fixture(scope="session")
-def alloydb_docker_compose_files() -> list[Path]:
-    return [Path(Path(__file__).parent / "docker-compose.alloydb-omni.yml")]
-
-
-@pytest.fixture(scope="session")
-def alloydb_docker_ip(alloydb_docker_services: DockerServiceRegistry) -> str:
-    return alloydb_docker_services.docker_ip
+if TYPE_CHECKING:
+    from collections.abc import Generator
 
 
-# alias to the latest
 @pytest.fixture(autouse=False, scope="session")
 def alloydb_omni_service(
-    alloydb_docker_services: DockerServiceRegistry,
-    default_alloydb_omni_service_name: str,
-    alloydb_docker_compose_files: list[Path],
-    alloydb_omni_port: int,
-    postgres_database: str,
-    postgres_user: str,
-    postgres_password: str,
-) -> Generator[None, None, None]:
-    os.environ["POSTGRES_PASSWORD"] = postgres_password
-    os.environ["POSTGRES_USER"] = postgres_user
-    os.environ["POSTGRES_DATABASE"] = postgres_database
-    os.environ[f"{default_alloydb_omni_service_name.upper()}_PORT"] = str(alloydb_omni_port)
-    alloydb_docker_services.start(
-        name=default_alloydb_omni_service_name,
-        docker_compose_files=alloydb_docker_compose_files,
-        timeout=45,
-        pause=1,
-        check=alloydb_omni_responsive,
-        port=alloydb_omni_port,
-        database=postgres_database,
-        user=postgres_user,
-        password=postgres_password,
-    )
-    yield
+    docker_service: DockerService,
+) -> Generator[DockerServiceRegistry, None, None]:
+    with _provide_postgres_service(
+        docker_service=docker_service, image="google/alloydbomni", name="alloydb-omni"
+    ) as service:
+        yield service
 
 
 @pytest.fixture(autouse=False, scope="session")
-def alloydb_omni_startup_connection(
-    alloydb_omni_service: DockerServiceRegistry,
-    alloydb_docker_ip: str,
-    alloydb_omni_port: int,
-    postgres_database: str,
-    postgres_user: str,
-    postgres_password: str,
-) -> Generator[psycopg.Connection, None, None]:
+def alloydb_omni_startup_connection(alloydb_omni_service: AlloyDBService) -> Generator[psycopg.Connection, None, None]:
     with psycopg.connect(
-        host=alloydb_docker_ip,
-        port=alloydb_omni_port,
-        user=postgres_user,
-        database=postgres_database,
-        password=postgres_password,
+        _make_connection_string(
+            host=alloydb_omni_service.host,
+            port=alloydb_omni_service.port,
+            user=alloydb_omni_service.user,
+            database=alloydb_omni_service.database,
+            password=alloydb_omni_service.password,
+        )
     ) as conn:
         yield conn

src/pytest_databases/docker/docker-compose.alloydb-omni.yml

@@ -21,7 +21,7 @@ def _make_connection_string(host: str, port: int, user: str, password: str, data
 
 @dataclasses.dataclass
 class PostgresService(ServiceContainer):
-    db: str
+    database: str
     password: str
     user: str
 
@@ -61,7 +61,7 @@ def check(_service: ServiceContainer) -> bool:
         exec_after_start=f"psql -U postgres -d postgres -c 'CREATE DATABASE {db_name};'",
     ) as service:
         yield PostgresService(
-            db=db_name,
+            database=db_name,
             host=service.host,
             port=service.port,
             user="postgres",
@@ -140,7 +140,7 @@ def postgres_startup_connection(
             port=postgres_service.port,
             user=postgres_service.user,
             password=postgres_service.password,
-            database=postgres_service.db,
+            database=postgres_service.database,
         ),
     ) as conn:
         yield conn
@@ -156,7 +156,7 @@ def postgres11_startup_connection(
             port=postgres_11_service.port,
             user=postgres_11_service.user,
             password=postgres_11_service.password,
-            database=postgres_11_service.db,
+            database=postgres_11_service.database,
         ),
     ) as conn:
         yield conn
@@ -172,7 +172,7 @@ def postgres12_startup_connection(
             port=postgres_12_service.port,
             user=postgres_12_service.user,
             password=postgres_12_service.password,
-            database=postgres_12_service.db,
+            database=postgres_12_service.database,
         ),
     ) as conn:
         yield conn
@@ -188,7 +188,7 @@ def postgres13_startup_connection(
             port=postgres_13_service.port,
             user=postgres_13_service.user,
             password=postgres_13_service.password,
-            database=postgres_13_service.db,
+            database=postgres_13_service.database,
         ),
     ) as conn:
         yield conn
@@ -204,7 +204,7 @@ def postgres14_startup_connection(
             port=postgres_14_service.port,
             user=postgres_14_service.user,
             password=postgres_14_service.password,
-            database=postgres_14_service.db,
+            database=postgres_14_service.database,
         ),
     ) as conn:
         yield conn
@@ -220,7 +220,7 @@ def postgres15_startup_connection(
             port=postgres_15_service.port,
             user=postgres_15_service.user,
             password=postgres_15_service.password,
-            database=postgres_15_service.db,
+            database=postgres_15_service.database,
         ),
     ) as conn:
         yield conn
@@ -236,7 +236,7 @@ def postgres16_startup_connection(
             port=postgres_16_service.port,
             user=postgres_16_service.user,
             password=postgres_16_service.password,
-            database=postgres_16_service.db,
+            database=postgres_16_service.database,
         ),
     ) as conn:
         yield conn
@@ -252,7 +252,7 @@ def postgres17_startup_connection(
             port=postgres_17_service.port,
             user=postgres_17_service.user,
             password=postgres_17_service.password,
-            database=postgres_17_service.db,
+            database=postgres_17_service.database,
         ),
     ) as conn:
         yield conn

src/pytest_databases/docker/postgres.py

@@ -2,58 +2,41 @@
 
 from typing import TYPE_CHECKING
 
-import pytest
+import psycopg
 
-from pytest_databases.docker.alloydb_omni import alloydb_omni_responsive
+from pytest_databases.docker.alloydb_omni import AlloyDBService
+from pytest_databases.docker.postgres import _make_connection_string
 
 if TYPE_CHECKING:
-    import psycopg
+    pass
 
-    from pytest_databases.docker import DockerServiceRegistry
 
-
-pytestmark = pytest.mark.skip
+# pytestmark = pytest.mark.skip
 
 pytest_plugins = [
     "pytest_databases.docker.alloydb_omni",
 ]
 
 
-def test_alloydb_omni_default_config(
-    default_alloydb_omni_service_name: str,
-    alloydb_omni_port: int,
-    postgres_database: str,
-    postgres_user: str,
-    postgres_password: str,
-) -> None:
-    assert default_alloydb_omni_service_name == "alloydb"
-    assert alloydb_omni_port == 5420
-    assert postgres_database == "postgres"
-    assert postgres_user == "postgres"
-    assert postgres_password == "super-secret"
-
-
-def test_alloydb_omni_services(
-    alloydb_docker_ip: str,
-    alloydb_omni_service: DockerServiceRegistry,
-    alloydb_omni_port: int,
-    postgres_database: str,
-    postgres_user: str,
-    postgres_password: str,
-) -> None:
-    ping = alloydb_omni_responsive(
-        alloydb_docker_ip,
-        port=alloydb_omni_port,
-        database=postgres_database,
-        user=postgres_user,
-        password=postgres_password,
-    )
-    assert ping
-
-
-def test_alloydb_omni_service_after_start(
-    postgres_startup_connection: psycopg.Connection,
-) -> None:
-    postgres_startup_connection.execute("CREATE TABLE if not exists simple_table as SELECT 1")
-    result = postgres_startup_connection.execute("select * from simple_table").fetchone()
+def check(service: AlloyDBService) -> bool:
+    with psycopg.connect(
+        _make_connection_string(
+            host=service.host,
+            port=service.port,
+            user=service.user,
+            database=service.database,
+            password=service.password,
+        )
+    ) as conn:
+        db_open = conn.execute("SELECT 1").fetchone()
+        return bool(db_open is not None and db_open[0] == 1)
+
+
+def test_alloydb_omni_services(alloydb_omni_service: AlloyDBService) -> None:
+    assert check(alloydb_omni_service)
+
+
+def test_alloydb_omni_service_after_start(alloydb_omni_startup_connection: psycopg.Connection) -> None:
+    alloydb_omni_startup_connection.execute("CREATE TABLE if not exists simple_table as SELECT 1")
+    result = alloydb_omni_startup_connection.execute("select * from simple_table").fetchone()
     assert bool(result is not None and result[0] == 1)