pytest-neon 2.2.0__py3-none-any.whl → 2.2.2__py3-none-any.whl

pytest_neon/__init__.py

@@ -9,7 +9,7 @@ from pytest_neon.plugin import (
     neon_engine,
 )
 
-__version__ = "2.2.0"
+__version__ = "2.2.2"
 __all__ = [
     "NeonBranch",
     "neon_branch",

pytest_neon/plugin.py

@@ -36,17 +36,19 @@ For full documentation, see: https://github.com/ZainRizvi/pytest-neon
 from __future__ import annotations
 
 import contextlib
+import json
 import os
 import random
 import time
 import warnings
 from collections.abc import Callable, Generator
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from datetime import datetime, timedelta, timezone
 from typing import Any, TypeVar
 
 import pytest
 import requests
+from filelock import FileLock
 from neon_api import NeonAPI
 from neon_api.exceptions import NeonAPIError
 from neon_api.schema import EndpointState
@@ -806,9 +808,24 @@ def _wait_for_operations(
         )
 
 
+def _branch_to_dict(branch: NeonBranch) -> dict[str, Any]:
+    """Convert NeonBranch to a JSON-serializable dict."""
+    return asdict(branch)
+
+
+def _dict_to_branch(data: dict[str, Any]) -> NeonBranch:
+    """Convert a dict back to NeonBranch."""
+    return NeonBranch(**data)
+
+
+# Timeout for waiting for migrations to complete (seconds)
+_MIGRATION_WAIT_TIMEOUT = 300  # 5 minutes
+
+
 @pytest.fixture(scope="session")
 def _neon_migration_branch(
     request: pytest.FixtureRequest,
+    tmp_path_factory: pytest.TempPathFactory,
 ) -> Generator[NeonBranch, None, None]:
     """
     Session-scoped branch where migrations are applied.
@@ -817,6 +834,13 @@ def _neon_migration_branch(
     the parent for all test branches. Migrations run once per session
     on this branch.
 
+    pytest-xdist Support:
+        When running with pytest-xdist, the first worker to acquire the lock
+        creates the migration branch. Other workers wait for migrations to
+        complete, then reuse the same branch. This avoids redundant API calls
+        and ensures migrations only run once. Only the creator cleans up the
+        branch at session end.
+
     Note: The migration branch cannot have an expiry because Neon doesn't
     allow creating child branches from branches with expiration dates.
     Cleanup relies on the fixture teardown at session end.
@@ -826,25 +850,123 @@ def _neon_migration_branch(
         it on request.config. After migrations run, _neon_branch_for_reset
         compares the fingerprint to detect if the schema actually changed.
     """
-    # No expiry - Neon doesn't allow children from branches with expiry
-    branch_gen = _create_neon_branch(
-        request,
-        branch_expiry_override=0,
-        branch_name_suffix="-migrated",
+    config = request.config
+    worker_id = _get_xdist_worker_id()
+    is_xdist = worker_id != "main"
+
+    # Get env var name for DATABASE_URL
+    env_var_name = _get_config_value(
+        config, "neon_env_var", "", "neon_env_var", "DATABASE_URL"
     )
-    branch = next(branch_gen)
 
-    # Capture schema fingerprint BEFORE migrations run
-    # This is stored on config so _neon_branch_for_reset can compare after
-    pre_migration_fingerprint = _get_schema_fingerprint(branch.connection_string)
-    request.config._neon_pre_migration_fingerprint = pre_migration_fingerprint  # type: ignore[attr-defined]
+    # For xdist, use shared temp directory and filelock
+    # tmp_path_factory.getbasetemp().parent is shared across all workers
+    if is_xdist:
+        root_tmp_dir = tmp_path_factory.getbasetemp().parent
+        cache_file = root_tmp_dir / "neon_migration_branch.json"
+        lock_file = root_tmp_dir / "neon_migration_branch.lock"
+        migrations_done_file = root_tmp_dir / "neon_migrations_done"
+    else:
+        cache_file = None
+        lock_file = None
+        migrations_done_file = None
+
+    is_creator = False
+    branch: NeonBranch
+    branch_gen: Generator[NeonBranch, None, None] | None = None
+    original_env_value: str | None = None
+
+    if is_xdist:
+        assert cache_file is not None and lock_file is not None
+        assert migrations_done_file is not None
+        with FileLock(str(lock_file)):
+            if cache_file.exists():
+                # Another worker already created the branch - reuse it
+                data = json.loads(cache_file.read_text())
+                branch = _dict_to_branch(data["branch"])
+                pre_migration_fingerprint = tuple(
+                    tuple(row) for row in data["pre_migration_fingerprint"]
+                )
+                config._neon_pre_migration_fingerprint = pre_migration_fingerprint  # type: ignore[attr-defined]
+
+                # Set DATABASE_URL for this worker (not done by _create_neon_branch)
+                original_env_value = os.environ.get(env_var_name)
+                os.environ[env_var_name] = branch.connection_string
+            else:
+                # First worker - create branch and cache it
+                is_creator = True
+                branch_gen = _create_neon_branch(
+                    request,
+                    branch_expiry_override=0,
+                    branch_name_suffix="-migrated",
+                )
+                branch = next(branch_gen)
+
+                # Capture schema fingerprint BEFORE migrations run
+                pre_migration_fingerprint = _get_schema_fingerprint(
+                    branch.connection_string
+                )
+                config._neon_pre_migration_fingerprint = pre_migration_fingerprint  # type: ignore[attr-defined]
+
+                # Cache for other workers (they'll read this after lock released)
+                # Note: We cache now with pre-migration fingerprint. The branch
+                # content will have migrations applied by neon_apply_migrations.
+                cache_file.write_text(
+                    json.dumps(
+                        {
+                            "branch": _branch_to_dict(branch),
+                            "pre_migration_fingerprint": pre_migration_fingerprint,
+                        }
+                    )
+                )
+
+        # Non-creator workers must wait for migrations to complete BEFORE
+        # neon_apply_migrations runs, otherwise they'll try to run migrations
+        # concurrently on the same branch, causing race conditions.
+        if not is_creator:
+            waited = 0.0
+            poll_interval = 0.5
+            while not migrations_done_file.exists():
+                if waited >= _MIGRATION_WAIT_TIMEOUT:
+                    raise RuntimeError(
+                        f"Timeout waiting for migrations to complete after "
+                        f"{_MIGRATION_WAIT_TIMEOUT}s. The creator worker may have "
+                        f"failed or is still running migrations."
+                    )
+                time.sleep(poll_interval)
+                waited += poll_interval
+    else:
+        # Not using xdist - create branch normally
+        is_creator = True
+        branch_gen = _create_neon_branch(
+            request,
+            branch_expiry_override=0,
+            branch_name_suffix="-migrated",
+        )
+        branch = next(branch_gen)
+
+        # Capture schema fingerprint BEFORE migrations run
+        pre_migration_fingerprint = _get_schema_fingerprint(branch.connection_string)
+        config._neon_pre_migration_fingerprint = pre_migration_fingerprint  # type: ignore[attr-defined]
+
+    # Mark whether this worker is the creator (used by neon_apply_migrations)
+    config._neon_is_migration_creator = is_creator  # type: ignore[attr-defined]
+    # Store migrations_done_file path for signaling after migrations complete
+    config._neon_migrations_done_file = migrations_done_file  # type: ignore[attr-defined]
 
     try:
         yield branch
     finally:
-        # Clean up by exhausting the generator (triggers branch deletion)
-        with contextlib.suppress(StopIteration):
-            next(branch_gen)
+        # Restore env var if we set it (non-creator workers)
+        if original_env_value is not None:
+            os.environ[env_var_name] = original_env_value
+        elif not is_creator and env_var_name in os.environ:
+            os.environ.pop(env_var_name, None)
+
+        # Only the creator cleans up the branch
+        if is_creator and branch_gen is not None:
+            with contextlib.suppress(StopIteration):
+                next(branch_gen)
 
 
 @pytest.fixture(scope="session")
@@ -855,6 +977,12 @@ def neon_apply_migrations(_neon_migration_branch: NeonBranch) -> Any:
     The migration branch is already created and DATABASE_URL is set.
     Migrations run once per test session, before any tests execute.
 
+    pytest-xdist Support:
+        When running with pytest-xdist, migrations only run on the first
+        worker (the one that created the migration branch). Other workers
+        wait for migrations to complete before proceeding. This ensures
+        migrations run exactly once, even with parallel workers.
+
     Smart Migration Detection:
         The plugin automatically detects whether migrations actually modified
         the database schema. If no schema changes occurred (or this fixture
@@ -897,11 +1025,41 @@ def neon_apply_migrations(_neon_migration_branch: NeonBranch) -> Any:
     return _MIGRATIONS_NOT_DEFINED
 
 
+@pytest.fixture(scope="session")
+def _neon_migrations_synchronized(
+    request: pytest.FixtureRequest,
+    _neon_migration_branch: NeonBranch,
+    neon_apply_migrations: Any,
+) -> Any:
+    """
+    Internal fixture that synchronizes migrations across xdist workers.
+
+    This fixture ensures that:
+    1. Only the creator worker runs migrations (non-creators wait in
+       _neon_migration_branch BEFORE neon_apply_migrations runs)
+    2. Creator signals completion after migrations finish
+    3. The return value from neon_apply_migrations is preserved for detection
+
+    Without xdist, this is a simple passthrough.
+    """
+    config = request.config
+    is_creator = getattr(config, "_neon_is_migration_creator", True)
+    migrations_done_file = getattr(config, "_neon_migrations_done_file", None)
+
+    if is_creator and migrations_done_file is not None:
+        # Creator: migrations just ran via neon_apply_migrations dependency
+        # Signal completion to other workers (who are waiting in
+        # _neon_migration_branch)
+        migrations_done_file.write_text("done")
+
+    return neon_apply_migrations
+
+
 @pytest.fixture(scope="session")
 def _neon_branch_for_reset(
     request: pytest.FixtureRequest,
     _neon_migration_branch: NeonBranch,
-    neon_apply_migrations: Any,  # Ensures migrations run first; value for detection
+    _neon_migrations_synchronized: Any,  # Ensures migrations complete; for detection
 ) -> Generator[NeonBranch, None, None]:
     """
     Internal fixture that creates a test branch from the migration branch.
@@ -931,7 +1089,8 @@ def _neon_branch_for_reset(
         - Migrations exist but are already applied (no schema changes)
     """
     # Check if migrations fixture was overridden
-    migrations_defined = neon_apply_migrations is not _MIGRATIONS_NOT_DEFINED
+    # _neon_migrations_synchronized passes through the neon_apply_migrations value
+    migrations_defined = _neon_migrations_synchronized is not _MIGRATIONS_NOT_DEFINED
 
     # Check if schema actually changed (if we have a pre-migration fingerprint)
     pre_fingerprint = getattr(request.config, "_neon_pre_migration_fingerprint", ())
@@ -1123,7 +1282,7 @@ def neon_branch(
 def neon_branch_shared(
     request: pytest.FixtureRequest,
     _neon_migration_branch: NeonBranch,
-    neon_apply_migrations: None,  # Ensures migrations run first
+    _neon_migrations_synchronized: Any,  # Ensures migrations complete first
 ) -> Generator[NeonBranch, None, None]:
     """
     Provide a shared Neon database branch for all tests in a module.

{pytest_neon-2.2.0.dist-info → pytest_neon-2.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytest-neon
-Version: 2.2.0
+Version: 2.2.2
 Summary: Pytest plugin for Neon database branch isolation in tests
 Project-URL: Homepage, https://github.com/ZainRizvi/pytest-neon
 Project-URL: Repository, https://github.com/ZainRizvi/pytest-neon
@@ -24,6 +24,7 @@ Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Database
 Classifier: Topic :: Software Development :: Testing
 Requires-Python: >=3.9
+Requires-Dist: filelock>=3.0
 Requires-Dist: neon-api>=0.1.0
 Requires-Dist: pytest>=7.0
 Requires-Dist: requests>=2.20

pytest_neon-2.2.2.dist-info/RECORD

@@ -0,0 +1,8 @@
+pytest_neon/__init__.py,sha256=NYxWUqqq-U4rhtxH4fhgZ7F4-vaxQ51oYKhuMFGhTes,398
+pytest_neon/plugin.py,sha256=8TciFbqE7mH86URTxSxw4sU8yqth17Mkrt2XOo0DZ-c,56023
+pytest_neon/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pytest_neon-2.2.2.dist-info/METADATA,sha256=Ga8VA-3kMmdhEj-oAer-qnOXkmZgC7GiBygh5HysiR4,19266
+pytest_neon-2.2.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+pytest_neon-2.2.2.dist-info/entry_points.txt,sha256=5U88Idj_G8-PSDb9VF3OwYFbGLHnGOo_GxgYvi0dtXw,37
+pytest_neon-2.2.2.dist-info/licenses/LICENSE,sha256=aKKp_Ex4WBHTByY4BhXJ181dzB_qYhi2pCUmZ7Spn_0,1067
+pytest_neon-2.2.2.dist-info/RECORD,,

pytest_neon-2.2.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-pytest_neon/__init__.py,sha256=-HXVGp38c6Ddmq90i5eo6OFxKVzBo3igJ6R3DKA8A7M,398
-pytest_neon/plugin.py,sha256=8y5LYvQossHJlPpJdh0DdZj8tFcO8OY45MMMfQytcRA,49269
-pytest_neon/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pytest_neon-2.2.0.dist-info/METADATA,sha256=dMgATSjnaUH1gH1posidEa99b_IE30ikeVsRsYC3oE0,19237
-pytest_neon-2.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-pytest_neon-2.2.0.dist-info/entry_points.txt,sha256=5U88Idj_G8-PSDb9VF3OwYFbGLHnGOo_GxgYvi0dtXw,37
-pytest_neon-2.2.0.dist-info/licenses/LICENSE,sha256=aKKp_Ex4WBHTByY4BhXJ181dzB_qYhi2pCUmZ7Spn_0,1067
-pytest_neon-2.2.0.dist-info/RECORD,,