Files
sqlmodel-pg-kit/tests/test_pg_integration.py
2025-08-17 22:18:45 +08:00

71 lines
2.2 KiB
Python

import os
import socket
import time
from typing import Optional
import pytest
from sqlmodel import SQLModel, Field
def _wait_for_port(host: str, port: int, deadline_s: float = 5.0) -> bool:
start = time.time()
while time.time() - start < deadline_s:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.settimeout(1.0)
try:
s.connect((host, port))
return True
except OSError:
time.sleep(0.25)
return False
@pytest.mark.integration
def test_postgres_connection_and_crud(monkeypatch):
# Prefer env to locate DB; fall back to localhost:5433
host = os.getenv("SQL_HOST") or os.getenv("PGHOST") or "127.0.0.1"
port = int(os.getenv("SQL_PORT") or os.getenv("PGPORT") or 5433)
if not _wait_for_port(host, port):
pytest.skip(f"Postgres not reachable at {host}:{port}; skip integration test")
from sqlmodel_pg_kit import db
from sqlmodel_pg_kit import create_all, Repository
from sqlmodel_pg_kit.db import get_session
class Hero(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
name: str
age: Optional[int] = None
user = os.getenv("SQL_USER") or os.getenv("PGUSER") or "postgres"
password = os.getenv("SQL_PASSWORD") or os.getenv("PGPASSWORD") or "change-me-strong"
database = os.getenv("SQL_DATABASE") or os.getenv("PGDATABASE") or "postgres"
sslmode = os.getenv("SQL_SSLMODE") or "disable"
cfg = db.DatabaseConfig(
host=host,
port=port,
user=user,
password=password,
database=database,
sslmode=sslmode,
echo=False,
)
monkeypatch.setattr(db, "cfg", cfg)
monkeypatch.setattr(db, "engine", db.create_engine(cfg.sync_url(), echo=False))
# Try a few retries in case container just started
for _ in range(10):
try:
create_all()
break
except Exception: # pragma: no cover - only on slow CI
time.sleep(0.5)
repo = Repository(Hero)
with get_session() as s:
repo.create(s, {"name": "PG Hero", "age": 1})
names = [h.name for h in repo.list(s)]
assert "PG Hero" in names