|
| 1 | +""" |
| 2 | +E2E test fixtures with real PostgreSQL database. |
| 3 | +
|
| 4 | +Uses a separate 'test_e2e' schema to isolate test data from production. |
| 5 | +Tests are skipped if DATABASE_URL is not set. |
| 6 | +
|
| 7 | +Note: These tests must NOT be run with pytest-xdist parallelization |
| 8 | +as multiple workers would conflict on the shared test_e2e schema. |
| 9 | +""" |
| 10 | + |
| 11 | +import os |
| 12 | + |
| 13 | +import pytest |
| 14 | +import pytest_asyncio |
| 15 | +from sqlalchemy import text |
| 16 | +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine |
| 17 | + |
| 18 | +from core.database.models import Base |
| 19 | + |
| 20 | + |
| 21 | +TEST_SCHEMA = "test_e2e" |
| 22 | + |
| 23 | +# Test data constants |
| 24 | +TEST_IMAGE_URL = "https://storage.googleapis.com/pyplots-images/test/plot.png" |
| 25 | +TEST_THUMB_URL = "https://storage.googleapis.com/pyplots-images/test/thumb.png" |
| 26 | + |
| 27 | + |
| 28 | +def _get_database_url(): |
| 29 | + """Get DATABASE_URL from environment, loading .env if needed.""" |
| 30 | + from dotenv import load_dotenv |
| 31 | + |
| 32 | + load_dotenv() |
| 33 | + return os.environ.get("DATABASE_URL") |
| 34 | + |
| 35 | + |
| 36 | +@pytest_asyncio.fixture(scope="function") |
| 37 | +async def pg_engine(): |
| 38 | + """ |
| 39 | + Create PostgreSQL engine and setup test schema. |
| 40 | +
|
| 41 | + Creates a separate 'test_e2e' schema to isolate tests from production data. |
| 42 | + The schema is dropped and recreated for each test. |
| 43 | + """ |
| 44 | + database_url = _get_database_url() |
| 45 | + if not database_url: |
| 46 | + pytest.skip("DATABASE_URL not set - skipping PostgreSQL E2E tests") |
| 47 | + |
| 48 | + # First create schema with a temporary engine (no search_path yet) |
| 49 | + temp_engine = create_async_engine(database_url, echo=False) |
| 50 | + async with temp_engine.begin() as conn: |
| 51 | + await conn.execute(text(f"DROP SCHEMA IF EXISTS {TEST_SCHEMA} CASCADE")) |
| 52 | + await conn.execute(text(f"CREATE SCHEMA {TEST_SCHEMA}")) |
| 53 | + await temp_engine.dispose() |
| 54 | + |
| 55 | + # Create engine with search_path set at connection level (handles pooling correctly) |
| 56 | + engine = create_async_engine( |
| 57 | + database_url, echo=False, connect_args={"server_settings": {"search_path": TEST_SCHEMA}} |
| 58 | + ) |
| 59 | + |
| 60 | + # Create tables in test schema |
| 61 | + async with engine.begin() as conn: |
| 62 | + await conn.run_sync(Base.metadata.create_all) |
| 63 | + |
| 64 | + yield engine |
| 65 | + |
| 66 | + # Cleanup: Drop entire test schema |
| 67 | + async with engine.begin() as conn: |
| 68 | + await conn.execute(text(f"DROP SCHEMA IF EXISTS {TEST_SCHEMA} CASCADE")) |
| 69 | + await engine.dispose() |
| 70 | + |
| 71 | + |
| 72 | +@pytest_asyncio.fixture |
| 73 | +async def pg_session(pg_engine): |
| 74 | + """Create session with test schema (search_path set at engine level).""" |
| 75 | + async_session = async_sessionmaker(pg_engine, class_=AsyncSession, expire_on_commit=False) |
| 76 | + async with async_session() as session: |
| 77 | + yield session |
| 78 | + await session.rollback() |
| 79 | + |
| 80 | + |
| 81 | +@pytest_asyncio.fixture |
| 82 | +async def pg_db_with_data(pg_session): |
| 83 | + """ |
| 84 | + Seed test schema with sample data. |
| 85 | +
|
| 86 | + Creates the same test data as tests/conftest.py:test_db_with_data |
| 87 | + but in the PostgreSQL test schema. |
| 88 | + """ |
| 89 | + from core.database.models import Impl, Library, Spec |
| 90 | + |
| 91 | + # Create libraries |
| 92 | + matplotlib_lib = Library( |
| 93 | + id="matplotlib", |
| 94 | + name="Matplotlib", |
| 95 | + version="3.10.0", |
| 96 | + documentation_url="https://matplotlib.org", |
| 97 | + description="Comprehensive library for visualizations", |
| 98 | + ) |
| 99 | + seaborn_lib = Library( |
| 100 | + id="seaborn", |
| 101 | + name="Seaborn", |
| 102 | + version="0.13.0", |
| 103 | + documentation_url="https://seaborn.pydata.org", |
| 104 | + description="Statistical data visualization", |
| 105 | + ) |
| 106 | + pg_session.add_all([matplotlib_lib, seaborn_lib]) |
| 107 | + |
| 108 | + # Create specs |
| 109 | + scatter_spec = Spec( |
| 110 | + id="scatter-basic", |
| 111 | + title="Basic Scatter Plot", |
| 112 | + description="A basic scatter plot", |
| 113 | + applications=["data visualization", "correlation analysis"], |
| 114 | + data=["numeric"], |
| 115 | + notes=["Use for 2D data"], |
| 116 | + tags={"plot_type": ["scatter"], "domain": ["statistics"], "data_type": ["numeric"], "features": ["basic"]}, |
| 117 | + issue=42, |
| 118 | + suggested="contributor", |
| 119 | + ) |
| 120 | + bar_spec = Spec( |
| 121 | + id="bar-grouped", |
| 122 | + title="Grouped Bar Chart", |
| 123 | + description="A grouped bar chart", |
| 124 | + applications=["comparisons"], |
| 125 | + data=["categorical", "numeric"], |
| 126 | + notes=["Good for comparing categories"], |
| 127 | + tags={"plot_type": ["bar"], "domain": ["statistics"], "data_type": ["categorical"], "features": ["grouped"]}, |
| 128 | + issue=43, |
| 129 | + suggested="contributor2", |
| 130 | + ) |
| 131 | + pg_session.add_all([scatter_spec, bar_spec]) |
| 132 | + await pg_session.commit() |
| 133 | + |
| 134 | + # Create implementations |
| 135 | + scatter_matplotlib = Impl( |
| 136 | + spec_id="scatter-basic", |
| 137 | + library_id="matplotlib", |
| 138 | + code="import matplotlib.pyplot as plt\n# scatter plot code", |
| 139 | + preview_url=TEST_IMAGE_URL.replace("plot", "scatter-matplotlib"), |
| 140 | + preview_thumb=TEST_THUMB_URL.replace("thumb", "scatter-matplotlib-thumb"), |
| 141 | + quality_score=92.5, |
| 142 | + generated_by="claude", |
| 143 | + python_version="3.13", |
| 144 | + library_version="3.10.0", |
| 145 | + ) |
| 146 | + scatter_seaborn = Impl( |
| 147 | + spec_id="scatter-basic", |
| 148 | + library_id="seaborn", |
| 149 | + code="import seaborn as sns\n# scatter plot code", |
| 150 | + preview_url=TEST_IMAGE_URL.replace("plot", "scatter-seaborn"), |
| 151 | + preview_thumb=TEST_THUMB_URL.replace("thumb", "scatter-seaborn-thumb"), |
| 152 | + quality_score=95.0, |
| 153 | + generated_by="claude", |
| 154 | + python_version="3.13", |
| 155 | + library_version="0.13.0", |
| 156 | + ) |
| 157 | + bar_matplotlib = Impl( |
| 158 | + spec_id="bar-grouped", |
| 159 | + library_id="matplotlib", |
| 160 | + code="import matplotlib.pyplot as plt\n# bar chart code", |
| 161 | + preview_url=TEST_IMAGE_URL.replace("plot", "bar-matplotlib"), |
| 162 | + quality_score=88.0, |
| 163 | + generated_by="claude", |
| 164 | + python_version="3.13", |
| 165 | + library_version="3.10.0", |
| 166 | + ) |
| 167 | + pg_session.add_all([scatter_matplotlib, scatter_seaborn, bar_matplotlib]) |
| 168 | + await pg_session.commit() |
| 169 | + |
| 170 | + # Expire all cached objects to ensure fresh loading with relationships |
| 171 | + pg_session.expire_all() |
| 172 | + |
| 173 | + yield pg_session |
0 commit comments