pytest-neon 2.1.0__py3-none-any.whl → 2.1.2__py3-none-any.whl

pytest_neon/__init__.py

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

pytest_neon/plugin.py

@@ -396,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
@@ -406,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",
@@ -416,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
@@ -434,6 +453,72 @@ 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
+    first_poll = True
+    while pending_op_ids and waited < max_wait_seconds:
+        # Poll immediately first time (operation usually completes instantly),
+        # then wait between subsequent polls
+        if first_poll:
+            time.sleep(0.1)  # Tiny delay to let operation start
+            waited += 0.1
+            first_poll = False
+        else:
+            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,

{pytest_neon-2.1.0.dist-info → pytest_neon-2.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytest-neon
-Version: 2.1.0
+Version: 2.1.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

pytest_neon-2.1.2.dist-info/RECORD

@@ -0,0 +1,8 @@
+pytest_neon/__init__.py,sha256=EKni9NpcUiIPy8J8vj-s6oKH6S98JlNs5JYuF151cMA,398
+pytest_neon/plugin.py,sha256=r30wz96rQZerAiAF8wAVmWPksF6lM2dsJyUNj2X9rWo,38038
+pytest_neon/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pytest_neon-2.1.2.dist-info/METADATA,sha256=QA7xaOCbi_wD8LTwb8gWQpZe8rG8BWcMrETS78HUOks,18734
+pytest_neon-2.1.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+pytest_neon-2.1.2.dist-info/entry_points.txt,sha256=5U88Idj_G8-PSDb9VF3OwYFbGLHnGOo_GxgYvi0dtXw,37
+pytest_neon-2.1.2.dist-info/licenses/LICENSE,sha256=aKKp_Ex4WBHTByY4BhXJ181dzB_qYhi2pCUmZ7Spn_0,1067
+pytest_neon-2.1.2.dist-info/RECORD,,

pytest_neon-2.1.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-pytest_neon/__init__.py,sha256=69SoGL2ciE38uvrXpEQPXiBFDkbWct7hl_k9MmswcSU,398
-pytest_neon/plugin.py,sha256=rb9CgOqle-sQHHoyFO5xDZ5wzrHx_4BLygiVywEYGbM,34949
-pytest_neon/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pytest_neon-2.1.0.dist-info/METADATA,sha256=ebdenaxT8v4GQeepZCglp04mb92iwJxp1PMnnS-9Nqs,18734
-pytest_neon-2.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-pytest_neon-2.1.0.dist-info/entry_points.txt,sha256=5U88Idj_G8-PSDb9VF3OwYFbGLHnGOo_GxgYvi0dtXw,37
-pytest_neon-2.1.0.dist-info/licenses/LICENSE,sha256=aKKp_Ex4WBHTByY4BhXJ181dzB_qYhi2pCUmZ7Spn_0,1067
-pytest_neon-2.1.0.dist-info/RECORD,,