Add real DB integration tests with testcontainers

12 tests covering all Queries methods against a real PostgreSQL instance via testcontainers. Automatically skipped when Docker is unavailable. Tests: api key lookup, last_used update, organization fetch, subscription CRUD, usage increment, cumulative increments.
2026-05-04 18:25:19 +00:00 · 2026-04-22 20:02:28 -05:00 · 2026-04-22 20:02:28 -05:00 · a62f1ecd03
commit a62f1ecd03
parent 3e16d44912
3 changed files with 501 additions and 0 deletions
--- a/packages/codeflash-api/pyproject.toml
+++ b/packages/codeflash-api/pyproject.toml
@ -27,6 +27,7 @@ dev = [
    "mypy>=1.14",
    "pytest>=9.0.3",
    "pytest-asyncio>=1.3.0",
+    "testcontainers[postgres]>=4.12",
 ]

 [tool.mypy]
--- a/packages/codeflash-api/tests/test_db_integration.py
+++ b/packages/codeflash-api/tests/test_db_integration.py
@ -0,0 +1,449 @@
+"""Integration tests for Queries against a real PostgreSQL instance.
+
+These tests require Docker.  They are automatically skipped when Docker
+is not available (e.g. local dev without Docker Desktop running).
+"""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+from datetime import datetime, timezone
+
+import asyncpg  # type: ignore[import-untyped]
+import pytest
+
+from codeflash_api.db._queries import Queries
+
+_DOCKER_AVAILABLE: bool | None = None
+
+
+def _is_docker_available() -> bool:
+    global _DOCKER_AVAILABLE  # noqa: PLW0603
+    if _DOCKER_AVAILABLE is None:
+        docker = shutil.which("docker")
+        if docker is None:
+            _DOCKER_AVAILABLE = False
+        else:
+            try:
+                subprocess.run(  # noqa: S603
+                    [docker, "info"],
+                    capture_output=True,
+                    timeout=5,
+                    check=True,
+                )
+                _DOCKER_AVAILABLE = True
+            except (
+                subprocess.CalledProcessError,
+                FileNotFoundError,
+                OSError,
+            ):
+                _DOCKER_AVAILABLE = False
+    return _DOCKER_AVAILABLE
+
+
+pytestmark = [
+    pytest.mark.skipif(
+        not _is_docker_available(),
+        reason="Docker is not available",
+    ),
+    pytest.mark.asyncio(),
+]
+
+SCHEMA_SQL = """\
+CREATE TABLE IF NOT EXISTS cf_api_keys (
+    id SERIAL PRIMARY KEY,
+    key TEXT NOT NULL,
+    suffix TEXT NOT NULL DEFAULT '',
+    name TEXT NOT NULL DEFAULT '',
+    user_id TEXT,
+    tier TEXT,
+    organization_id TEXT,
+    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+    last_used TIMESTAMPTZ
+);
+
+CREATE TABLE IF NOT EXISTS organizations (
+    id TEXT PRIMARY KEY,
+    name TEXT NOT NULL,
+    github_org_id TEXT,
+    description TEXT,
+    website TEXT,
+    auto_add_github_members BOOLEAN NOT NULL DEFAULT TRUE,
+    added_by TEXT,
+    privacy_mode BOOLEAN NOT NULL DEFAULT FALSE,
+    subscription BOOLEAN NOT NULL DEFAULT FALSE,
+    created_at TIMESTAMPTZ DEFAULT NOW(),
+    updated_at TIMESTAMPTZ DEFAULT NOW()
+);
+
+CREATE TABLE IF NOT EXISTS subscriptions (
+    id TEXT PRIMARY KEY,
+    user_id TEXT NOT NULL UNIQUE,
+    plan_type TEXT NOT NULL,
+    optimizations_used INT NOT NULL DEFAULT 0,
+    total_lifetime_optimizations INT NOT NULL DEFAULT 0,
+    optimizations_limit INT NOT NULL DEFAULT 4000,
+    subscription_status TEXT NOT NULL DEFAULT 'active',
+    stripe_customer_id TEXT,
+    stripe_subscription_id TEXT,
+    current_period_start TIMESTAMPTZ,
+    current_period_end TIMESTAMPTZ,
+    cancel_at_period_end BOOLEAN NOT NULL DEFAULT FALSE,
+    cancellation_request_date TIMESTAMPTZ,
+    created_at TIMESTAMPTZ DEFAULT NOW(),
+    updated_at TIMESTAMPTZ DEFAULT NOW()
+);
+"""
+
+
+@pytest.fixture(scope="module")
+def pg_container():
+    """
+    Start a PostgreSQL container for the test module.
+    """
+    from testcontainers.postgres import PostgresContainer  # type: ignore[import-untyped]  # noqa: I001
+
+    with PostgresContainer("postgres:16-alpine") as pg:
+        yield pg
+
+
+@pytest.fixture(scope="module")
+def pg_dsn(pg_container) -> str:  # type: ignore[no-untyped-def]
+    """
+    Connection string for the test Postgres instance.
+    """
+    return pg_container.get_connection_url().replace(
+        "postgresql+psycopg2://", "postgresql://"
+    )
+
+
+@pytest.fixture
+async def pool(pg_dsn: str) -> asyncpg.Pool:
+    """
+    Create a fresh pool and apply the schema for each test.
+    """
+    p = await asyncpg.create_pool(pg_dsn, min_size=1, max_size=5)
+    assert p is not None
+    async with p.acquire() as conn:
+        await conn.execute(SCHEMA_SQL)
+        await conn.execute("DELETE FROM subscriptions")
+        await conn.execute("DELETE FROM cf_api_keys")
+        await conn.execute("DELETE FROM organizations")
+    yield p
+    await p.close()
+
+
+@pytest.fixture
+def queries(pool: asyncpg.Pool) -> Queries:
+    """
+    Queries instance backed by the test pool.
+    """
+    return Queries(pool=pool)
+
+
+async def _insert_api_key(
+    pool: asyncpg.Pool,
+    *,
+    key: str = "hashed_abc",
+    user_id: str = "user-1",
+    suffix: str = "xabc",
+    name: str = "Test Key",
+) -> int:
+    """
+    Insert a test API key row and return its id.
+    """
+    row = await pool.fetchrow(
+        "INSERT INTO cf_api_keys"
+        " (key, suffix, name, user_id)"
+        " VALUES ($1, $2, $3, $4)"
+        " RETURNING id",
+        key,
+        suffix,
+        name,
+        user_id,
+    )
+    assert row is not None
+    return int(row["id"])
+
+
+class TestGetApiKeyByHash:
+    """Tests for Queries.get_api_key_by_hash against real Postgres."""
+
+    async def test_returns_matching_key(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        A key inserted with a known hash is found by that hash.
+        """
+        await _insert_api_key(pool, key="sha384_hash_abc")
+
+        result = await queries.get_api_key_by_hash("sha384_hash_abc")
+
+        assert result is not None
+        assert "user-1" == result.user_id
+        assert "xabc" == result.suffix
+
+    async def test_returns_none_for_missing_hash(
+        self,
+        queries: Queries,
+    ) -> None:
+        """
+        A hash with no matching row returns None.
+        """
+        result = await queries.get_api_key_by_hash("nonexistent_hash")
+
+        assert result is None
+
+    async def test_returns_correct_fields(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        All APIKey fields are populated from the database row.
+        """
+        await pool.execute(
+            "INSERT INTO cf_api_keys"
+            " (key, suffix, name, user_id, tier, organization_id)"
+            " VALUES ($1, $2, $3, $4, $5, $6)",
+            "full_fields_hash",
+            "zz99",
+            "Full Key",
+            "user-42",
+            "PRO",
+            "org-7",
+        )
+
+        result = await queries.get_api_key_by_hash("full_fields_hash")
+
+        assert result is not None
+        assert "full_fields_hash" == result.key
+        assert "zz99" == result.suffix
+        assert "Full Key" == result.name
+        assert "user-42" == result.user_id
+        assert "PRO" == result.tier
+        assert "org-7" == result.organization_id
+        assert isinstance(result.created_at, datetime)
+
+
+class TestUpdateApiKeyLastUsed:
+    """Tests for Queries.update_api_key_last_used against real Postgres."""
+
+    async def test_updates_timestamp(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        Calling update sets last_used to a recent timestamp.
+        """
+        key_id = await _insert_api_key(pool, key="update_test_hash")
+
+        before = await pool.fetchval(
+            "SELECT last_used FROM cf_api_keys WHERE id = $1",
+            key_id,
+        )
+        assert before is None
+
+        await queries.update_api_key_last_used(key_id)
+
+        after = await pool.fetchval(
+            "SELECT last_used FROM cf_api_keys WHERE id = $1",
+            key_id,
+        )
+        assert after is not None
+        assert isinstance(after, datetime)
+
+
+class TestGetOrganization:
+    """Tests for Queries.get_organization against real Postgres."""
+
+    async def test_returns_matching_org(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        An inserted organization is found by ID.
+        """
+        await pool.execute(
+            "INSERT INTO organizations (id, name, privacy_mode)"
+            " VALUES ($1, $2, $3)",
+            "org-1",
+            "Acme Corp",
+            True,
+        )
+
+        result = await queries.get_organization("org-1")
+
+        assert result is not None
+        assert "Acme Corp" == result.name
+        assert result.privacy_mode is True
+
+    async def test_returns_none_for_missing_org(
+        self,
+        queries: Queries,
+    ) -> None:
+        """
+        A missing organization ID returns None.
+        """
+        result = await queries.get_organization("nonexistent-org")
+
+        assert result is None
+
+
+class TestGetSubscription:
+    """Tests for Queries.get_subscription against real Postgres."""
+
+    async def test_returns_matching_subscription(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        An inserted subscription is found by user_id.
+        """
+        await pool.execute(
+            "INSERT INTO subscriptions"
+            " (id, user_id, plan_type, optimizations_limit)"
+            " VALUES ($1, $2, $3, $4)",
+            "sub-1",
+            "user-1",
+            "PRO",
+            100_000,
+        )
+
+        result = await queries.get_subscription("user-1")
+
+        assert result is not None
+        assert "PRO" == result.plan_type
+        assert 100_000 == result.optimizations_limit
+        assert 0 == result.optimizations_used
+
+    async def test_returns_none_for_missing_user(
+        self,
+        queries: Queries,
+    ) -> None:
+        """
+        A user_id with no subscription returns None.
+        """
+        result = await queries.get_subscription("ghost-user")
+
+        assert result is None
+
+
+class TestCreateSubscription:
+    """Tests for Queries.create_subscription against real Postgres."""
+
+    async def test_creates_and_returns_subscription(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        A new subscription is persisted and returned with defaults.
+        """
+        result = await queries.create_subscription(
+            user_id="new-user",
+            plan_type="FREE",
+            optimizations_limit=4000,
+        )
+
+        assert "new-user" == result.user_id
+        assert "FREE" == result.plan_type
+        assert 4000 == result.optimizations_limit
+        assert 0 == result.optimizations_used
+
+        row = await pool.fetchrow(
+            "SELECT * FROM subscriptions WHERE user_id = $1",
+            "new-user",
+        )
+        assert row is not None
+        assert "active" == row["subscription_status"]
+
+    async def test_string_limit_is_cast_to_int(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        A string optimizations_limit is cast to int before insert.
+        """
+        result = await queries.create_subscription(
+            user_id="str-limit-user",
+            plan_type="PRO",
+            optimizations_limit="100000",
+        )
+
+        assert 100_000 == result.optimizations_limit
+
+        db_val = await pool.fetchval(
+            "SELECT optimizations_limit FROM subscriptions"
+            " WHERE user_id = $1",
+            "str-limit-user",
+        )
+        assert 100_000 == db_val
+
+
+class TestIncrementUsage:
+    """Tests for Queries.increment_usage against real Postgres."""
+
+    async def test_increments_both_counters(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        Both optimizations_used and total_lifetime_optimizations increase.
+        """
+        await pool.execute(
+            "INSERT INTO subscriptions"
+            " (id, user_id, plan_type)"
+            " VALUES ($1, $2, $3)",
+            "sub-inc",
+            "inc-user",
+            "FREE",
+        )
+
+        await queries.increment_usage("inc-user", cost=5)
+
+        row = await pool.fetchrow(
+            "SELECT optimizations_used, total_lifetime_optimizations"
+            " FROM subscriptions WHERE user_id = $1",
+            "inc-user",
+        )
+        assert row is not None
+        assert 5 == row["optimizations_used"]
+        assert 5 == row["total_lifetime_optimizations"]
+
+    async def test_increments_are_cumulative(
+        self,
+        queries: Queries,
+        pool: asyncpg.Pool,
+    ) -> None:
+        """
+        Multiple increments accumulate correctly.
+        """
+        await pool.execute(
+            "INSERT INTO subscriptions"
+            " (id, user_id, plan_type)"
+            " VALUES ($1, $2, $3)",
+            "sub-cum",
+            "cum-user",
+            "FREE",
+        )
+
+        await queries.increment_usage("cum-user", cost=3)
+        await queries.increment_usage("cum-user", cost=7)
+
+        row = await pool.fetchrow(
+            "SELECT optimizations_used, total_lifetime_optimizations"
+            " FROM subscriptions WHERE user_id = $1",
+            "cum-user",
+        )
+        assert row is not None
+        assert 10 == row["optimizations_used"]
+        assert 10 == row["total_lifetime_optimizations"]
--- a/uv.lock
+++ b/uv.lock
@ -501,8 +501,10 @@ dependencies = [
 [package.dev-dependencies]
 dev = [
    { name = "httpx" },
+    { name = "mypy" },
    { name = "pytest" },
    { name = "pytest-asyncio" },
+    { name = "testcontainers" },
 ]

 [package.metadata]
@ -524,8 +526,10 @@ requires-dist = [
 [package.metadata.requires-dev]
 dev = [
    { name = "httpx", specifier = ">=0.28.1" },
+    { name = "mypy", specifier = ">=1.14" },
    { name = "pytest", specifier = ">=9.0.3" },
    { name = "pytest-asyncio", specifier = ">=1.3.0" },
+    { name = "testcontainers", extras = ["postgres"], specifier = ">=4.12" },
 ]

 [[package]]
@ -1023,6 +1027,20 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" },
 ]

+[[package]]
+name = "docker"
+version = "7.1.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pywin32", marker = "sys_platform == 'win32'" },
+    { name = "requests" },
+    { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/91/9b/4a2ea29aeba62471211598dac5d96825bb49348fa07e906ea930394a83ce/docker-7.1.0.tar.gz", hash = "sha256:ad8c70e6e3f8926cb8a92619b832b4ea5299e2831c14284663184e200546fa6c", size = 117834, upload-time = "2024-05-23T11:13:57.216Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e3/26/57c6fb270950d476074c087527a558ccb6f4436657314bfb6cdf484114c4/docker-7.1.0-py3-none-any.whl", hash = "sha256:c96b93b7f0a746f9e77d325bcfb87422a3d8bd4f03136ae8a85b37f1898d5fc0", size = 147774, upload-time = "2024-05-23T11:13:55.01Z" },
+]
+
 [[package]]
 name = "docstring-parser"
 version = "0.18.0"
@ -1437,6 +1455,7 @@ dependencies = [
    { name = "opt-einsum" },
    { name = "scipy" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/50/f0/bcb81d28267d2054d0daed766c7fa16bcee5e481331b4d1e14f5fbe662be/jax-0.10.0.tar.gz", hash = "sha256:0119c767de1645f407df72345d28a3837dc904f1d698911c121d8f2b396fdece", size = 2663397, upload-time = "2026-04-22T13:22:28.563Z" }
 wheels = [
    { url = "https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl", hash = "sha256:76c42ba163c8db3dc2e449e225b888c0edfb623ded31efdc96d85e0fda1d26e8", size = 3094950, upload-time = "2026-04-16T12:32:11.576Z" },
 ]
@ -3107,6 +3126,22 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" },
 ]

+[[package]]
+name = "pywin32"
+version = "311"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e7/ab/01ea1943d4eba0f850c3c61e78e8dd59757ff815ff3ccd0a84de5f541f42/pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31", size = 8706543, upload-time = "2025-07-14T20:13:20.765Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/a8/a0e8d07d4d051ec7502cd58b291ec98dcc0c3fff027caad0470b72cfcc2f/pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067", size = 9495040, upload-time = "2025-07-14T20:13:22.543Z" },
+    { url = "https://files.pythonhosted.org/packages/ba/3a/2ae996277b4b50f17d61f0603efd8253cb2d79cc7ae159468007b586396d/pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852", size = 8710102, upload-time = "2025-07-14T20:13:24.682Z" },
+    { url = "https://files.pythonhosted.org/packages/a5/be/3fd5de0979fcb3994bfee0d65ed8ca9506a8a1260651b86174f6a86f52b3/pywin32-311-cp313-cp313-win32.whl", hash = "sha256:f95ba5a847cba10dd8c4d8fefa9f2a6cf283b8b88ed6178fa8a6c1ab16054d0d", size = 8705700, upload-time = "2025-07-14T20:13:26.471Z" },
+    { url = "https://files.pythonhosted.org/packages/e3/28/e0a1909523c6890208295a29e05c2adb2126364e289826c0a8bc7297bd5c/pywin32-311-cp313-cp313-win_amd64.whl", hash = "sha256:718a38f7e5b058e76aee1c56ddd06908116d35147e133427e59a3983f703a20d", size = 9494700, upload-time = "2025-07-14T20:13:28.243Z" },
+    { url = "https://files.pythonhosted.org/packages/04/bf/90339ac0f55726dce7d794e6d79a18a91265bdf3aa70b6b9ca52f35e022a/pywin32-311-cp313-cp313-win_arm64.whl", hash = "sha256:7b4075d959648406202d92a2310cb990fea19b535c7f4a78d3f5e10b926eeb8a", size = 8709318, upload-time = "2025-07-14T20:13:30.348Z" },
+    { url = "https://files.pythonhosted.org/packages/c9/31/097f2e132c4f16d99a22bfb777e0fd88bd8e1c634304e102f313af69ace5/pywin32-311-cp314-cp314-win32.whl", hash = "sha256:b7a2c10b93f8986666d0c803ee19b5990885872a7de910fc460f9b0c2fbf92ee", size = 8840714, upload-time = "2025-07-14T20:13:32.449Z" },
+    { url = "https://files.pythonhosted.org/packages/90/4b/07c77d8ba0e01349358082713400435347df8426208171ce297da32c313d/pywin32-311-cp314-cp314-win_amd64.whl", hash = "sha256:3aca44c046bd2ed8c90de9cb8427f581c479e594e99b5c0bb19b29c10fd6cb87", size = 9656800, upload-time = "2025-07-14T20:13:34.312Z" },
+    { url = "https://files.pythonhosted.org/packages/c0/d2/21af5c535501a7233e734b8af901574572da66fcc254cb35d0609c9080dd/pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42", size = 8932540, upload-time = "2025-07-14T20:13:36.379Z" },
+]
+
 [[package]]
 name = "pyyaml"
 version = "6.0.3"
@ -3467,6 +3502,22 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl", hash = "sha256:cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5", size = 7734, upload-time = "2025-12-29T12:55:20.718Z" },
 ]

+[[package]]
+name = "testcontainers"
+version = "4.14.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "docker" },
+    { name = "python-dotenv" },
+    { name = "typing-extensions" },
+    { name = "urllib3" },
+    { name = "wrapt" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ca/ac/a597c3a0e02b26cbed6dd07df68be1e57684766fd1c381dee9b170a99690/testcontainers-4.14.2.tar.gz", hash = "sha256:1340ccf16fe3acd9389a6c9e1d9ab21d9fe99a8afdf8165f89c3e69c1967d239", size = 166841, upload-time = "2026-03-18T05:19:16.696Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/13/2d/26b8b30067d94339afee62c3edc9b803a6eb9332f521ba77d8aaab5de873/testcontainers-4.14.2-py3-none-any.whl", hash = "sha256:0d0522c3cd8f8d9627cda41f7a6b51b639fa57bdc492923c045117933c668d68", size = 125712, upload-time = "2026-03-18T05:19:15.29Z" },
+]
+
 [[package]]
 name = "textual"
 version = "8.2.4"