pytest-neon 2.0.0__tar.gz → 2.1.1__tar.gz

{pytest_neon-2.0.0 → pytest_neon-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytest-neon
-Version: 2.0.0
+Version: 2.1.1
 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
@@ -489,16 +489,26 @@ The `neon_branch_readwrite` fixture uses Neon's branch restore API to reset data
 
 This is similar to database transactions but at the branch level.
 
-## Limitations
+## Parallel Test Execution (pytest-xdist)
 
-### Parallel Test Execution
+This plugin supports parallel test execution with [pytest-xdist](https://pytest-xdist.readthedocs.io/). Each xdist worker automatically gets its own isolated branch.
 
-This plugin sets the `DATABASE_URL` environment variable, which is process-global. This means it is **not compatible with pytest-xdist** or other parallel test runners that run tests in the same process.
+```bash
+# Run tests in parallel with 4 workers
+pip install pytest-xdist
+pytest -n 4
+```
+
+**How it works:**
+- Each xdist worker (gw0, gw1, gw2, etc.) creates its own branch
+- Branches are named with the worker ID suffix (e.g., `-test-gw0`, `-test-gw1`)
+- Workers run tests in parallel without database state interference
+- All branches are cleaned up after the test session
 
-If you need parallel execution, you can:
-- Use `neon_branch.connection_string` directly instead of relying on `DATABASE_URL`
-- Run with `pytest-xdist --dist=loadfile` to keep modules in separate processes
-- Run tests serially (default pytest behavior)
+**Cost implications:**
+- Running with `-n 4` creates 4 branches (one per worker) plus the migration branch
+- Choose your parallelism level based on your Neon plan's branch limits
+- Each worker's branch is reset after each test using the fast reset operation (~0.5s)
 
 ## Troubleshooting
 

{pytest_neon-2.0.0 → pytest_neon-2.1.1}/README.md

@@ -444,16 +444,26 @@ The `neon_branch_readwrite` fixture uses Neon's branch restore API to reset data
 
 This is similar to database transactions but at the branch level.
 
-## Limitations
+## Parallel Test Execution (pytest-xdist)
 
-### Parallel Test Execution
+This plugin supports parallel test execution with [pytest-xdist](https://pytest-xdist.readthedocs.io/). Each xdist worker automatically gets its own isolated branch.
 
-This plugin sets the `DATABASE_URL` environment variable, which is process-global. This means it is **not compatible with pytest-xdist** or other parallel test runners that run tests in the same process.
+```bash
+# Run tests in parallel with 4 workers
+pip install pytest-xdist
+pytest -n 4
+```
+
+**How it works:**
+- Each xdist worker (gw0, gw1, gw2, etc.) creates its own branch
+- Branches are named with the worker ID suffix (e.g., `-test-gw0`, `-test-gw1`)
+- Workers run tests in parallel without database state interference
+- All branches are cleaned up after the test session
 
-If you need parallel execution, you can:
-- Use `neon_branch.connection_string` directly instead of relying on `DATABASE_URL`
-- Run with `pytest-xdist --dist=loadfile` to keep modules in separate processes
-- Run tests serially (default pytest behavior)
+**Cost implications:**
+- Running with `-n 4` creates 4 branches (one per worker) plus the migration branch
+- Choose your parallelism level based on your Neon plan's branch limits
+- Each worker's branch is reset after each test using the fast reset operation (~0.5s)
 
 ## Troubleshooting
 

{pytest_neon-2.0.0 → pytest_neon-2.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pytest-neon"
-version = "2.0.0"
+version = "2.1.1"
 description = "Pytest plugin for Neon database branch isolation in tests"
 readme = "README.md"
 license = "MIT"

{pytest_neon-2.0.0 → pytest_neon-2.1.1}/src/pytest_neon/__init__.py

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

{pytest_neon-2.0.0 → pytest_neon-2.1.1}/src/pytest_neon/plugin.py

@@ -56,6 +56,17 @@ DEFAULT_BRANCH_EXPIRY_SECONDS = 600
 _MIGRATIONS_NOT_DEFINED = object()
 
 
+def _get_xdist_worker_id() -> str:
+    """
+    Get the pytest-xdist worker ID, or "main" if not running under xdist.
+
+    When running tests in parallel with pytest-xdist, each worker process
+    gets a unique ID (gw0, gw1, gw2, etc.). This is used to create separate
+    branches per worker to avoid database state pollution between parallel tests.
+    """
+    return os.environ.get("PYTEST_XDIST_WORKER", "main")
+
+
 def _get_schema_fingerprint(connection_string: str) -> tuple[tuple[Any, ...], ...]:
     """
     Get a fingerprint of the database schema for change detection.
@@ -385,7 +396,8 @@ def _reset_branch_to_parent(
     """Reset a branch to its parent's state using the Neon API.
 
     Uses exponential backoff retry logic to handle transient API errors
-    that can occur during parallel test execution.
+    that can occur during parallel test execution. After initiating the
+    restore, polls the operation status until it completes.
 
     Args:
         branch: The branch to reset
@@ -395,7 +407,10 @@ def _reset_branch_to_parent(
     if not branch.parent_id:
         raise RuntimeError(f"Branch {branch.branch_id} has no parent - cannot reset")
 
-    url = f"https://console.neon.tech/api/v2/projects/{branch.project_id}/branches/{branch.branch_id}/restore"
+    base_url = "https://console.neon.tech/api/v2"
+    project_id = branch.project_id
+    branch_id = branch.branch_id
+    restore_url = f"{base_url}/projects/{project_id}/branches/{branch_id}/restore"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json",
@@ -405,12 +420,27 @@ def _reset_branch_to_parent(
     for attempt in range(max_retries + 1):
         try:
             response = requests.post(
-                url,
+                restore_url,
                 headers=headers,
                 json={"source_branch_id": branch.parent_id},
                 timeout=30,
             )
             response.raise_for_status()
+
+            # The restore API returns operations that run asynchronously.
+            # We must wait for operations to complete before the next test
+            # starts, otherwise connections may fail during the restore.
+            data = response.json()
+            operations = data.get("operations", [])
+
+            if operations:
+                _wait_for_operations(
+                    project_id=branch.project_id,
+                    operations=operations,
+                    headers=headers,
+                    base_url=base_url,
+                )
+
             return  # Success
         except requests.RequestException as e:
             last_error = e
@@ -423,6 +453,64 @@ def _reset_branch_to_parent(
     raise last_error  # type: ignore[misc]
 
 
+def _wait_for_operations(
+    project_id: str,
+    operations: list[dict[str, Any]],
+    headers: dict[str, str],
+    base_url: str,
+    max_wait_seconds: float = 60,
+    poll_interval: float = 0.5,
+) -> None:
+    """Wait for Neon operations to complete.
+
+    Args:
+        project_id: The Neon project ID
+        operations: List of operation dicts from the API response
+        headers: HTTP headers including auth
+        base_url: Base URL for Neon API
+        max_wait_seconds: Maximum time to wait (default: 60s)
+        poll_interval: Time between polls (default: 0.5s)
+    """
+    # Get operation IDs that aren't already finished
+    pending_op_ids = [
+        op["id"] for op in operations if op.get("status") not in ("finished", "skipped")
+    ]
+
+    if not pending_op_ids:
+        return  # All operations already complete
+
+    waited = 0.0
+    while pending_op_ids and waited < max_wait_seconds:
+        time.sleep(poll_interval)
+        waited += poll_interval
+
+        # Check status of each pending operation
+        still_pending = []
+        for op_id in pending_op_ids:
+            op_url = f"{base_url}/projects/{project_id}/operations/{op_id}"
+            try:
+                response = requests.get(op_url, headers=headers, timeout=10)
+                response.raise_for_status()
+                op_data = response.json().get("operation", {})
+                status = op_data.get("status")
+
+                if status == "error":
+                    err = op_data.get("error", "unknown error")
+                    raise RuntimeError(f"Operation {op_id} failed: {err}")
+                if status not in ("finished", "skipped"):
+                    still_pending.append(op_id)
+            except requests.RequestException:
+                # On network error, assume still pending and retry
+                still_pending.append(op_id)
+
+        pending_op_ids = still_pending
+
+    if pending_op_ids:
+        raise RuntimeError(
+            f"Timeout waiting for operations to complete: {pending_op_ids}"
+        )
+
+
 @pytest.fixture(scope="session")
 def _neon_migration_branch(
     request: pytest.FixtureRequest,
@@ -527,6 +615,12 @@ def _neon_branch_for_reset(
     session, avoiding issues with Python's module caching (e.g., SQLAlchemy
     engines created at import time would otherwise point to stale branches).
 
+    Parallel Test Support (pytest-xdist):
+        When running tests in parallel with pytest-xdist, each worker gets its
+        own branch. This prevents database state pollution between tests running
+        concurrently on different workers. The worker ID is included in the
+        branch name suffix (e.g., "-test-gw0", "-test-gw1").
+
     Smart Migration Detection:
         This fixture implements a cost-optimization strategy:
 
@@ -558,15 +652,21 @@ def _neon_branch_for_reset(
         # Assume migrations changed something to be safe
         schema_changed = True
 
+    # Get worker ID for parallel test support
+    # Each xdist worker gets its own branch to avoid state pollution
+    worker_id = _get_xdist_worker_id()
+    branch_suffix = f"-test-{worker_id}"
+
     # Only create a child branch if migrations actually modified the schema
-    if schema_changed:
+    # OR if we're running under xdist (each worker needs its own branch)
+    if schema_changed or worker_id != "main":
         yield from _create_neon_branch(
             request,
             parent_branch_id_override=_neon_migration_branch.branch_id,
-            branch_name_suffix="-test",
+            branch_name_suffix=branch_suffix,
         )
     else:
-        # No schema changes - reuse the migration branch directly
+        # No schema changes and not parallel - reuse the migration branch directly
         # This saves creating an unnecessary branch
         yield _neon_migration_branch
 

{pytest_neon-2.0.0 → pytest_neon-2.1.1}/tests/test_reset_behavior.py

@@ -21,6 +21,8 @@ class TestResetRetryBehavior:
         # Mock requests.post to fail twice, then succeed
         mock_response = mocker.Mock()
         mock_response.raise_for_status = mocker.Mock()
+        # Return empty operations list (already complete)
+        mock_response.json.return_value = {"operations": []}
 
         import requests
 
@@ -103,6 +105,8 @@ class TestResetRetryBehavior:
 
         mock_response = mocker.Mock()
         mock_response.raise_for_status = mocker.Mock()
+        # Return empty operations list (already complete)
+        mock_response.json.return_value = {"operations": []}
         mock_post = mocker.patch(
             "pytest_neon.plugin.requests.post", return_value=mock_response
         )
@@ -113,6 +117,51 @@ class TestResetRetryBehavior:
         assert mock_post.call_count == 1
         assert mock_sleep.call_count == 0
 
+    def test_reset_waits_for_operations_to_complete(self, mocker):
+        """Verify reset polls operation status until complete."""
+        branch = NeonBranch(
+            branch_id="br-test",
+            project_id="proj-test",
+            connection_string="postgresql://test",
+            host="test.neon.tech",
+            parent_id="br-parent",
+        )
+
+        # Mock POST response with pending operation
+        mock_post_response = mocker.Mock()
+        mock_post_response.raise_for_status = mocker.Mock()
+        mock_post_response.json.return_value = {
+            "operations": [{"id": "op-123", "status": "running"}]
+        }
+
+        # Mock GET responses: first running, then finished
+        get_call_count = [0]
+
+        def mock_get(*args, **kwargs):
+            get_call_count[0] += 1
+            mock_get_response = mocker.Mock()
+            mock_get_response.raise_for_status = mocker.Mock()
+            if get_call_count[0] < 3:
+                mock_get_response.json.return_value = {
+                    "operation": {"id": "op-123", "status": "running"}
+                }
+            else:
+                mock_get_response.json.return_value = {
+                    "operation": {"id": "op-123", "status": "finished"}
+                }
+            return mock_get_response
+
+        mocker.patch(
+            "pytest_neon.plugin.requests.post", return_value=mock_post_response
+        )
+        mocker.patch("pytest_neon.plugin.requests.get", side_effect=mock_get)
+        mocker.patch("pytest_neon.plugin.time.sleep")
+
+        _reset_branch_to_parent(branch, "fake-api-key")
+
+        # Should have polled until finished
+        assert get_call_count[0] == 3
+
 
 class TestResetBehavior:
     """Test that branch reset happens between tests."""

pytest_neon-2.1.1/tests/test_xdist_worker_support.py

@@ -0,0 +1,189 @@
+"""Tests for pytest-xdist parallel worker support."""
+
+
+class TestXdistBranchIsolation:
+    """Test that parallel workers get separate branches."""
+
+    def test_xdist_worker_creates_branch_even_without_migrations(
+        self, pytester, monkeypatch
+    ):
+        """Even without schema changes, xdist workers should get their own branch."""
+        monkeypatch.setenv("PYTEST_XDIST_WORKER", "gw0")
+
+        pytester.makeconftest(
+            """
+            import os
+            import pytest
+            from pytest_neon.plugin import (
+                NeonBranch,
+                _get_xdist_worker_id,
+                _MIGRATIONS_NOT_DEFINED,
+            )
+
+            branch_creation_calls = []
+
+            @pytest.fixture(scope="session")
+            def _neon_migration_branch(request):
+                # Store empty fingerprint to simulate no psycopg available
+                request.config._neon_pre_migration_fingerprint = ()
+
+                return NeonBranch(
+                    branch_id="br-migration-123",
+                    project_id="proj-mock",
+                    connection_string="postgresql://mock:mock@migration.neon.tech/mockdb",
+                    host="migration.neon.tech",
+                    parent_id="br-parent-000",
+                )
+
+            @pytest.fixture(scope="session")
+            def neon_apply_migrations(_neon_migration_branch):
+                # Return sentinel to simulate NO migrations defined
+                return _MIGRATIONS_NOT_DEFINED
+
+            @pytest.fixture(scope="session")
+            def _neon_branch_for_reset(
+                request, _neon_migration_branch, neon_apply_migrations
+            ):
+                # Replicate the real logic
+                sentinel = _MIGRATIONS_NOT_DEFINED
+                migrations_defined = neon_apply_migrations is not sentinel
+                fingerprint_key = "_neon_pre_migration_fingerprint"
+                pre_fp = getattr(request.config, fingerprint_key, ())
+                schema_changed = False
+
+                if migrations_defined and pre_fp:
+                    schema_changed = False  # Simplified
+                elif migrations_defined and not pre_fp:
+                    schema_changed = True
+
+                worker_id = _get_xdist_worker_id()
+                suffix = f"-test-{worker_id}"
+
+                # Key: create branch even without schema changes when xdist
+                if schema_changed or worker_id != "main":
+                    branch_creation_calls.append(f"created-branch{suffix}")
+                    conn = f"postgresql://mock:mock@t{suffix}.neon.tech/db"
+                    branch_info = NeonBranch(
+                        branch_id=f"br-test{suffix}",
+                        project_id="proj-mock",
+                        connection_string=conn,
+                        host=f"t{suffix}.neon.tech",
+                        parent_id=_neon_migration_branch.branch_id,
+                    )
+                else:
+                    branch_creation_calls.append("reused-migration-branch")
+                    branch_info = _neon_migration_branch
+
+                os.environ["DATABASE_URL"] = branch_info.connection_string
+                try:
+                    yield branch_info
+                finally:
+                    os.environ.pop("DATABASE_URL", None)
+
+            @pytest.fixture(scope="session", autouse=True)
+            def verify_branch_created():
+                yield
+                # Under xdist, should create a new branch
+                assert branch_creation_calls == ["created-branch-test-gw0"]
+            """
+        )
+
+        pytester.makepyfile(
+            """
+            def test_xdist_creates_branch(_neon_branch_for_reset):
+                # Just trigger the fixture
+                assert _neon_branch_for_reset is not None
+            """
+        )
+
+        result = pytester.runpytest("-v")
+        result.assert_outcomes(passed=1)
+
+    def test_non_xdist_reuses_migration_branch_without_migrations(
+        self, pytester, monkeypatch
+    ):
+        """Without xdist and without migrations, should reuse migration branch."""
+        monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False)
+
+        pytester.makeconftest(
+            """
+            import os
+            import pytest
+            from pytest_neon.plugin import (
+                NeonBranch,
+                _get_xdist_worker_id,
+                _MIGRATIONS_NOT_DEFINED,
+            )
+
+            branch_creation_calls = []
+
+            @pytest.fixture(scope="session")
+            def _neon_migration_branch(request):
+                request.config._neon_pre_migration_fingerprint = ()
+                return NeonBranch(
+                    branch_id="br-migration-123",
+                    project_id="proj-mock",
+                    connection_string="postgresql://mock:mock@migration.neon.tech/mockdb",
+                    host="migration.neon.tech",
+                    parent_id="br-parent-000",
+                )
+
+            @pytest.fixture(scope="session")
+            def neon_apply_migrations(_neon_migration_branch):
+                return _MIGRATIONS_NOT_DEFINED
+
+            @pytest.fixture(scope="session")
+            def _neon_branch_for_reset(
+                request, _neon_migration_branch, neon_apply_migrations
+            ):
+                sentinel = _MIGRATIONS_NOT_DEFINED
+                migrations_defined = neon_apply_migrations is not sentinel
+                fingerprint_key = "_neon_pre_migration_fingerprint"
+                pre_fp = getattr(request.config, fingerprint_key, ())
+                schema_changed = False
+
+                if migrations_defined and pre_fp:
+                    schema_changed = False
+                elif migrations_defined and not pre_fp:
+                    schema_changed = True
+
+                worker_id = _get_xdist_worker_id()
+                suffix = f"-test-{worker_id}"
+
+                if schema_changed or worker_id != "main":
+                    branch_creation_calls.append(f"created-branch{suffix}")
+                    conn = f"postgresql://mock:mock@t{suffix}.neon.tech/db"
+                    branch_info = NeonBranch(
+                        branch_id=f"br-test{suffix}",
+                        project_id="proj-mock",
+                        connection_string=conn,
+                        host=f"t{suffix}.neon.tech",
+                        parent_id=_neon_migration_branch.branch_id,
+                    )
+                else:
+                    branch_creation_calls.append("reused-migration-branch")
+                    branch_info = _neon_migration_branch
+
+                os.environ["DATABASE_URL"] = branch_info.connection_string
+                try:
+                    yield branch_info
+                finally:
+                    os.environ.pop("DATABASE_URL", None)
+
+            @pytest.fixture(scope="session", autouse=True)
+            def verify_branch_reused():
+                yield
+                # Without xdist, should reuse migration branch
+                assert branch_creation_calls == ["reused-migration-branch"]
+            """
+        )
+
+        pytester.makepyfile(
+            """
+            def test_no_xdist_reuses_branch(_neon_branch_for_reset):
+                assert _neon_branch_for_reset is not None
+            """
+        )
+
+        result = pytester.runpytest("-v")
+        result.assert_outcomes(passed=1)

{pytest_neon-2.0.0 → pytest_neon-2.1.1}/uv.lock

@@ -1200,7 +1200,7 @@ wheels = [
 
 [[package]]
 name = "pytest-neon"
-version = "1.0.0"
+version = "2.0.0"
 source = { editable = "." }
 dependencies = [
     { name = "neon-api" },