|
| 1 | +""" |
| 2 | +Pytest configuration for Tortoise ORM tests. |
| 3 | +
|
| 4 | +Uses function-scoped fixtures for true test isolation. |
| 5 | +""" |
| 6 | + |
1 | 7 | import os |
2 | 8 |
|
3 | 9 | import pytest |
| 10 | +import pytest_asyncio |
4 | 11 |
|
5 | | -from tortoise.contrib.test import finalizer, initializer |
| 12 | +from tortoise.context import tortoise_test_context |
6 | 13 |
|
7 | 14 |
|
8 | 15 | @pytest.fixture(scope="session", autouse=True) |
9 | | -def initialize_tests(request): |
10 | | - # Reduce the default timeout for psycopg because the tests become very slow otherwise |
| 16 | +def configure_psycopg(): |
| 17 | + """Configure psycopg timeout for faster tests.""" |
11 | 18 | try: |
12 | 19 | from tortoise.backends.psycopg import PsycopgClient |
13 | 20 |
|
14 | 21 | PsycopgClient.default_timeout = float(os.environ.get("TORTOISE_POSTGRES_TIMEOUT", "15")) |
15 | 22 | except ImportError: |
16 | 23 | pass |
17 | 24 |
|
| 25 | + |
| 26 | +# ============================================================================ |
| 27 | +# HELPER FUNCTIONS |
| 28 | +# ============================================================================ |
| 29 | + |
| 30 | + |
| 31 | +async def _truncate_all_tables(ctx) -> None: |
| 32 | + """Truncate all tables in the given context.""" |
| 33 | + if ctx.apps: |
| 34 | + for model in ctx.apps.get_models_iterable(): |
| 35 | + quote_char = model._meta.db.query_class.SQL_CONTEXT.quote_char |
| 36 | + await model._meta.db.execute_script( |
| 37 | + f"DELETE FROM {quote_char}{model._meta.db_table}{quote_char}" # nosec |
| 38 | + ) |
| 39 | + |
| 40 | + |
| 41 | +# ============================================================================ |
| 42 | +# PYTEST FIXTURES FOR TESTS |
| 43 | +# These fixtures provide different isolation patterns for test scenarios |
| 44 | +# ============================================================================ |
| 45 | + |
| 46 | + |
| 47 | +@pytest_asyncio.fixture(scope="module") |
| 48 | +async def db_module(): |
| 49 | + """ |
| 50 | + Module-scoped fixture: Creates TortoiseContext once per test module. |
| 51 | +
|
| 52 | + This is the base fixture that creates the database schema once per module. |
| 53 | + Other fixtures build on top of this for different isolation strategies. |
| 54 | +
|
| 55 | + Note: Uses connection_label="models" to match standard test infrastructure. |
| 56 | + """ |
18 | 57 | db_url = os.getenv("TORTOISE_TEST_DB", "sqlite://:memory:") |
19 | | - initializer(["tests.testmodels"], db_url=db_url) |
20 | | - request.addfinalizer(finalizer) |
| 58 | + async with tortoise_test_context( |
| 59 | + modules=["tests.testmodels"], |
| 60 | + db_url=db_url, |
| 61 | + app_label="models", |
| 62 | + connection_label="models", |
| 63 | + ) as ctx: |
| 64 | + yield ctx |
| 65 | + |
| 66 | + |
| 67 | +@pytest_asyncio.fixture(scope="function") |
| 68 | +async def db(db_module): |
| 69 | + """ |
| 70 | + Function-scoped fixture with transaction rollback cleanup. |
| 71 | +
|
| 72 | + Each test runs inside a transaction that gets rolled back at the end, |
| 73 | + providing isolation without the overhead of schema recreation. |
| 74 | +
|
| 75 | + For databases that don't support transactions (e.g., MySQL MyISAM), |
| 76 | + falls back to truncation cleanup. |
| 77 | +
|
| 78 | + This is the FASTEST isolation method - use for most tests. |
| 79 | +
|
| 80 | + Usage: |
| 81 | + @pytest.mark.asyncio |
| 82 | + async def test_something(db): |
| 83 | + obj = await Model.create(name="test") |
| 84 | + assert obj.id is not None |
| 85 | + # Changes are rolled back after test |
| 86 | + """ |
| 87 | + # Get connection from the context using its default connection |
| 88 | + conn = db_module.db() |
| 89 | + |
| 90 | + # Check if the database supports transactions |
| 91 | + if conn.capabilities.supports_transactions: |
| 92 | + # Start a savepoint/transaction |
| 93 | + transaction = conn._in_transaction() |
| 94 | + await transaction.__aenter__() |
| 95 | + |
| 96 | + try: |
| 97 | + yield db_module |
| 98 | + finally: |
| 99 | + # Rollback the transaction (discards all changes made during test) |
| 100 | + class _RollbackException(Exception): |
| 101 | + pass |
| 102 | + |
| 103 | + await transaction.__aexit__(_RollbackException, _RollbackException(), None) |
| 104 | + else: |
| 105 | + # For databases without transaction support (e.g., MyISAM), |
| 106 | + # fall back to truncation cleanup |
| 107 | + yield db_module |
| 108 | + await _truncate_all_tables(db_module) |
| 109 | + |
| 110 | + |
| 111 | +@pytest_asyncio.fixture(scope="function") |
| 112 | +async def db_simple(db_module): |
| 113 | + """ |
| 114 | + Function-scoped fixture with NO cleanup between tests. |
| 115 | +
|
| 116 | + Tests share state - data from one test persists to the next within the module. |
| 117 | + Use ONLY for read-only tests or tests that manage their own cleanup. |
| 118 | +
|
| 119 | + Usage: |
| 120 | + @pytest.mark.asyncio |
| 121 | + async def test_read_only(db_simple): |
| 122 | + # Read-only operations, no writes |
| 123 | + config = get_config() |
| 124 | + assert "host" in config |
| 125 | + """ |
| 126 | + yield db_module |
| 127 | + |
| 128 | + |
| 129 | +@pytest_asyncio.fixture(scope="function") |
| 130 | +async def db_isolated(): |
| 131 | + """ |
| 132 | + Function-scoped fixture with full database recreation per test. |
| 133 | +
|
| 134 | + Creates a completely fresh database for EACH test. This is the SLOWEST |
| 135 | + method but provides maximum isolation. |
| 136 | +
|
| 137 | + Use when: |
| 138 | + - Testing database creation/dropping |
| 139 | + - Tests need custom model modules |
| 140 | + - Tests must have completely clean state |
| 141 | +
|
| 142 | + Usage: |
| 143 | + @pytest.mark.asyncio |
| 144 | + async def test_with_fresh_db(db_isolated): |
| 145 | + # Completely fresh database |
| 146 | + ... |
| 147 | + """ |
| 148 | + db_url = os.getenv("TORTOISE_TEST_DB", "sqlite://:memory:") |
| 149 | + async with tortoise_test_context( |
| 150 | + modules=["tests.testmodels"], |
| 151 | + db_url=db_url, |
| 152 | + app_label="models", |
| 153 | + connection_label="models", |
| 154 | + ) as ctx: |
| 155 | + yield ctx |
| 156 | + |
| 157 | + |
| 158 | +@pytest_asyncio.fixture(scope="function") |
| 159 | +async def db_truncate(db_module): |
| 160 | + """ |
| 161 | + Function-scoped fixture with table truncation cleanup. |
| 162 | +
|
| 163 | + After each test, all tables are truncated (DELETE FROM). |
| 164 | + Faster than db_isolated but slower than db (transaction rollback). |
| 165 | +
|
| 166 | + Use when testing transaction behavior (can't use rollback for cleanup). |
| 167 | +
|
| 168 | + Usage: |
| 169 | + @pytest.mark.asyncio |
| 170 | + async def test_with_transactions(db_truncate): |
| 171 | + async with in_transaction(): |
| 172 | + await Model.create(name="test") |
| 173 | + # Table truncated after test |
| 174 | + """ |
| 175 | + yield db_module |
| 176 | + await _truncate_all_tables(db_module) |
| 177 | + |
| 178 | + |
| 179 | +# ============================================================================ |
| 180 | +# HELPER FIXTURES |
| 181 | +# ============================================================================ |
| 182 | + |
| 183 | + |
| 184 | +def make_db_fixture( |
| 185 | + modules: list[str], app_label: str = "models", connection_label: str = "models" |
| 186 | +): |
| 187 | + """ |
| 188 | + Factory function to create custom db fixtures with different modules. |
| 189 | +
|
| 190 | + Use this in subdirectory conftest.py files for tests that need |
| 191 | + custom model modules. |
| 192 | +
|
| 193 | + Example usage in tests/fields/conftest.py: |
| 194 | + db_array = make_db_fixture(["tests.fields.test_array"]) |
| 195 | +
|
| 196 | + Args: |
| 197 | + modules: List of module paths to discover models from. |
| 198 | + app_label: The app label for the models, defaults to "models". |
| 199 | + connection_label: The connection alias name, defaults to "models". |
| 200 | +
|
| 201 | + Returns: |
| 202 | + An async fixture function. |
| 203 | + """ |
| 204 | + |
| 205 | + @pytest_asyncio.fixture(scope="function") |
| 206 | + async def _db_fixture(): |
| 207 | + db_url = os.getenv("TORTOISE_TEST_DB", "sqlite://:memory:") |
| 208 | + async with tortoise_test_context( |
| 209 | + modules=modules, |
| 210 | + db_url=db_url, |
| 211 | + app_label=app_label, |
| 212 | + connection_label=connection_label, |
| 213 | + ) as ctx: |
| 214 | + yield ctx |
| 215 | + |
| 216 | + return _db_fixture |
0 commit comments