pytest-neon 2.1.4__py3-none-any.whl → 2.2.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
pytest_neon/__init__.py CHANGED
@@ -9,7 +9,7 @@ from pytest_neon.plugin import (
9
9
  neon_engine,
10
10
  )
11
11
 
12
- __version__ = "2.1.4"
12
+ __version__ = "2.2.1"
13
13
  __all__ = [
14
14
  "NeonBranch",
15
15
  "neon_branch",
pytest_neon/plugin.py CHANGED
@@ -36,26 +36,196 @@ For full documentation, see: https://github.com/ZainRizvi/pytest-neon
36
36
  from __future__ import annotations
37
37
 
38
38
  import contextlib
39
+ import json
39
40
  import os
41
+ import random
40
42
  import time
41
43
  import warnings
42
- from collections.abc import Generator
43
- from dataclasses import dataclass
44
+ from collections.abc import Callable, Generator
45
+ from dataclasses import asdict, dataclass
44
46
  from datetime import datetime, timedelta, timezone
45
- from typing import Any
47
+ from typing import Any, TypeVar
46
48
 
47
49
  import pytest
48
50
  import requests
51
+ from filelock import FileLock
49
52
  from neon_api import NeonAPI
53
+ from neon_api.exceptions import NeonAPIError
50
54
  from neon_api.schema import EndpointState
51
55
 
56
+ T = TypeVar("T")
57
+
52
58
  # Default branch expiry in seconds (10 minutes)
53
59
  DEFAULT_BRANCH_EXPIRY_SECONDS = 600
54
60
 
61
+ # Rate limit retry configuration
62
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
63
+ # Neon limits: 700 requests/minute (~11/sec), burst up to 40/sec per route
64
+ _RATE_LIMIT_BASE_DELAY = 4.0 # seconds
65
+ _RATE_LIMIT_MAX_TOTAL_DELAY = 90.0 # 1.5 minutes total cap
66
+ _RATE_LIMIT_JITTER_FACTOR = 0.25 # +/- 25% jitter
67
+ _RATE_LIMIT_MAX_ATTEMPTS = 10 # Maximum number of retry attempts
68
+
55
69
  # Sentinel value to detect when neon_apply_migrations was not overridden
56
70
  _MIGRATIONS_NOT_DEFINED = object()
57
71
 
58
72
 
73
+ class NeonRateLimitError(Exception):
74
+ """Raised when Neon API rate limit is exceeded and retries are exhausted."""
75
+
76
+ pass
77
+
78
+
79
+ def _calculate_retry_delay(
80
+ attempt: int,
81
+ base_delay: float = _RATE_LIMIT_BASE_DELAY,
82
+ jitter_factor: float = _RATE_LIMIT_JITTER_FACTOR,
83
+ ) -> float:
84
+ """
85
+ Calculate delay for a retry attempt with exponential backoff and jitter.
86
+
87
+ Args:
88
+ attempt: The retry attempt number (0-indexed)
89
+ base_delay: Base delay in seconds
90
+ jitter_factor: Jitter factor (0.25 means +/- 25%)
91
+
92
+ Returns:
93
+ Delay in seconds with jitter applied
94
+ """
95
+ # Exponential backoff: base_delay * 2^attempt
96
+ delay = base_delay * (2**attempt)
97
+
98
+ # Apply jitter: delay * (1 +/- jitter_factor)
99
+ jitter = delay * jitter_factor * (2 * random.random() - 1)
100
+ return delay + jitter
101
+
102
+
103
+ def _is_rate_limit_error(exc: Exception) -> bool:
104
+ """
105
+ Check if an exception indicates a rate limit (429) error.
106
+
107
+ Handles both requests.HTTPError (with response object) and NeonAPIError
108
+ (which only has the error text, not the response object).
109
+
110
+ Args:
111
+ exc: The exception to check
112
+
113
+ Returns:
114
+ True if this is a rate limit error, False otherwise
115
+ """
116
+ # Check NeonAPIError first - it inherits from HTTPError but doesn't have
117
+ # a response object, so we need to check the error text
118
+ if isinstance(exc, NeonAPIError):
119
+ # NeonAPIError doesn't preserve the response object, only the text
120
+ # Check for rate limit indicators in the error message
121
+ # Note: We use "too many requests" specifically to avoid false positives
122
+ # from errors like "too many connections" or "too many rows"
123
+ error_text = str(exc).lower()
124
+ return (
125
+ "429" in error_text
126
+ or "rate limit" in error_text
127
+ or "too many requests" in error_text
128
+ )
129
+ if isinstance(exc, requests.HTTPError):
130
+ return exc.response is not None and exc.response.status_code == 429
131
+ return False
132
+
133
+
134
+ def _get_retry_after_from_error(exc: Exception) -> float | None:
135
+ """
136
+ Extract Retry-After header value from an exception if available.
137
+
138
+ Args:
139
+ exc: The exception to check
140
+
141
+ Returns:
142
+ The Retry-After value in seconds, or None if not available
143
+ """
144
+ if isinstance(exc, requests.HTTPError) and exc.response is not None:
145
+ retry_after = exc.response.headers.get("Retry-After")
146
+ if retry_after:
147
+ try:
148
+ return float(retry_after)
149
+ except ValueError:
150
+ pass
151
+ return None
152
+
153
+
154
+ def _retry_on_rate_limit(
155
+ operation: Callable[[], T],
156
+ operation_name: str,
157
+ base_delay: float = _RATE_LIMIT_BASE_DELAY,
158
+ max_total_delay: float = _RATE_LIMIT_MAX_TOTAL_DELAY,
159
+ jitter_factor: float = _RATE_LIMIT_JITTER_FACTOR,
160
+ max_attempts: int = _RATE_LIMIT_MAX_ATTEMPTS,
161
+ ) -> T:
162
+ """
163
+ Execute an operation with retry logic for rate limit (429) errors.
164
+
165
+ Uses exponential backoff with jitter. Retries until the operation succeeds,
166
+ the total delay exceeds max_total_delay, or max_attempts is reached.
167
+
168
+ See: https://api-docs.neon.tech/reference/api-rate-limiting
169
+
170
+ Args:
171
+ operation: Callable that may raise requests.HTTPError or NeonAPIError
172
+ operation_name: Human-readable name for error messages
173
+ base_delay: Base delay in seconds for first retry
174
+ max_total_delay: Maximum total delay across all retries
175
+ jitter_factor: Jitter factor for randomization
176
+ max_attempts: Maximum number of retry attempts
177
+
178
+ Returns:
179
+ The result of the operation
180
+
181
+ Raises:
182
+ NeonRateLimitError: If rate limit retries are exhausted
183
+ requests.HTTPError: For non-429 HTTP errors
184
+ NeonAPIError: For non-429 API errors
185
+ Exception: For other errors from the operation
186
+ """
187
+ total_delay = 0.0
188
+ attempt = 0
189
+
190
+ while True:
191
+ try:
192
+ return operation()
193
+ except (requests.HTTPError, NeonAPIError) as e:
194
+ if _is_rate_limit_error(e):
195
+ # Check for Retry-After header (may be added by Neon in future)
196
+ retry_after = _get_retry_after_from_error(e)
197
+ if retry_after is not None:
198
+ # Ensure minimum delay to prevent infinite loops if Retry-After is 0
199
+ delay = max(retry_after, 0.1)
200
+ else:
201
+ delay = _calculate_retry_delay(attempt, base_delay, jitter_factor)
202
+
203
+ # Check if we've exceeded max total delay
204
+ if total_delay + delay > max_total_delay:
205
+ raise NeonRateLimitError(
206
+ f"Rate limit exceeded for {operation_name}. "
207
+ f"Max total delay ({max_total_delay:.1f}s) reached after "
208
+ f"{attempt + 1} attempts. "
209
+ f"See: https://api-docs.neon.tech/reference/api-rate-limiting"
210
+ ) from e
211
+
212
+ # Check if we've exceeded max attempts
213
+ attempt += 1
214
+ if attempt >= max_attempts:
215
+ raise NeonRateLimitError(
216
+ f"Rate limit exceeded for {operation_name}. "
217
+ f"Max attempts ({max_attempts}) reached after "
218
+ f"{total_delay:.1f}s total delay. "
219
+ f"See: https://api-docs.neon.tech/reference/api-rate-limiting"
220
+ ) from e
221
+
222
+ time.sleep(delay)
223
+ total_delay += delay
224
+ else:
225
+ # Non-429 error, re-raise immediately
226
+ raise
227
+
228
+
59
229
  def _get_xdist_worker_id() -> str:
60
230
  """
61
231
  Get the pytest-xdist worker ID, or "main" if not running under xdist.
@@ -166,7 +336,12 @@ def _get_default_branch_id(neon: NeonAPI, project_id: str) -> str | None:
166
336
  The branch ID of the default branch, or None if not found.
167
337
  """
168
338
  try:
169
- response = neon.branches(project_id=project_id)
339
+ # Wrap in retry logic to handle rate limits
340
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
341
+ response = _retry_on_rate_limit(
342
+ lambda: neon.branches(project_id=project_id),
343
+ operation_name="list_branches",
344
+ )
170
345
  for branch in response.branches:
171
346
  # Check both 'default' and 'primary' flags for compatibility
172
347
  if getattr(branch, "default", False) or getattr(branch, "primary", False):
@@ -375,10 +550,15 @@ def _create_neon_branch(
375
550
  branch_config["expires_at"] = expires_at.strftime("%Y-%m-%dT%H:%M:%SZ")
376
551
 
377
552
  # Create branch with compute endpoint
378
- result = neon.branch_create(
379
- project_id=project_id,
380
- branch=branch_config,
381
- endpoints=[{"type": "read_write"}],
553
+ # Wrap in retry logic to handle rate limits
554
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
555
+ result = _retry_on_rate_limit(
556
+ lambda: neon.branch_create(
557
+ project_id=project_id,
558
+ branch=branch_config,
559
+ endpoints=[{"type": "read_write"}],
560
+ ),
561
+ operation_name="branch_create",
382
562
  )
383
563
 
384
564
  branch = result.branch
@@ -402,8 +582,11 @@ def _create_neon_branch(
402
582
  waited = 0.0
403
583
 
404
584
  while True:
405
- endpoint_response = neon.endpoint(
406
- project_id=project_id, endpoint_id=endpoint_id
585
+ # Wrap in retry logic to handle rate limits during polling
586
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
587
+ endpoint_response = _retry_on_rate_limit(
588
+ lambda: neon.endpoint(project_id=project_id, endpoint_id=endpoint_id),
589
+ operation_name="endpoint_status",
407
590
  )
408
591
  endpoint = endpoint_response.endpoint
409
592
  state = endpoint.current_state
@@ -436,10 +619,15 @@ def _create_neon_branch(
436
619
 
437
620
  # Reset password to get the password value
438
621
  # (newly created branches don't expose password)
439
- password_response = neon.role_password_reset(
440
- project_id=project_id,
441
- branch_id=branch.id,
442
- role_name=role_name,
622
+ # Wrap in retry logic to handle rate limits
623
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
624
+ password_response = _retry_on_rate_limit(
625
+ lambda: neon.role_password_reset(
626
+ project_id=project_id,
627
+ branch_id=branch.id,
628
+ role_name=role_name,
629
+ ),
630
+ operation_name="role_password_reset",
443
631
  )
444
632
  password = password_response.role.password
445
633
 
@@ -472,30 +660,34 @@ def _create_neon_branch(
472
660
  # Cleanup: delete branch unless --neon-keep-branches was specified
473
661
  if not keep_branches:
474
662
  try:
475
- neon.branch_delete(project_id=project_id, branch_id=branch.id)
663
+ # Wrap in retry logic to handle rate limits
664
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
665
+ _retry_on_rate_limit(
666
+ lambda: neon.branch_delete(
667
+ project_id=project_id, branch_id=branch.id
668
+ ),
669
+ operation_name="branch_delete",
670
+ )
476
671
  except Exception as e:
477
672
  # Log but don't fail tests due to cleanup issues
478
- import warnings
479
-
480
673
  warnings.warn(
481
674
  f"Failed to delete Neon branch {branch.id}: {e}",
482
675
  stacklevel=2,
483
676
  )
484
677
 
485
678
 
486
- def _reset_branch_to_parent(
487
- branch: NeonBranch, api_key: str, max_retries: int = 3
488
- ) -> None:
679
+ def _reset_branch_to_parent(branch: NeonBranch, api_key: str) -> None:
489
680
  """Reset a branch to its parent's state using the Neon API.
490
681
 
491
- Uses exponential backoff retry logic to handle transient API errors
492
- that can occur during parallel test execution. After initiating the
493
- restore, polls the operation status until it completes.
682
+ Uses exponential backoff retry logic with jitter to handle rate limit (429)
683
+ errors. After initiating the restore, polls the operation status until it
684
+ completes.
685
+
686
+ See: https://api-docs.neon.tech/reference/api-rate-limiting
494
687
 
495
688
  Args:
496
689
  branch: The branch to reset
497
690
  api_key: Neon API key
498
- max_retries: Maximum number of retry attempts (default: 3)
499
691
  """
500
692
  if not branch.parent_id:
501
693
  raise RuntimeError(f"Branch {branch.branch_id} has no parent - cannot reset")
@@ -509,41 +701,31 @@ def _reset_branch_to_parent(
509
701
  "Content-Type": "application/json",
510
702
  }
511
703
 
512
- last_error: Exception | None = None
513
- for attempt in range(max_retries + 1):
514
- try:
515
- response = requests.post(
516
- restore_url,
517
- headers=headers,
518
- json={"source_branch_id": branch.parent_id},
519
- timeout=30,
520
- )
521
- response.raise_for_status()
522
-
523
- # The restore API returns operations that run asynchronously.
524
- # We must wait for operations to complete before the next test
525
- # starts, otherwise connections may fail during the restore.
526
- data = response.json()
527
- operations = data.get("operations", [])
528
-
529
- if operations:
530
- _wait_for_operations(
531
- project_id=branch.project_id,
532
- operations=operations,
533
- headers=headers,
534
- base_url=base_url,
535
- )
536
-
537
- return # Success
538
- except requests.RequestException as e:
539
- last_error = e
540
- if attempt < max_retries:
541
- # Exponential backoff: 1s, 2s, 4s
542
- wait_time = 2**attempt
543
- time.sleep(wait_time)
544
-
545
- # All retries exhausted
546
- raise last_error # type: ignore[misc]
704
+ def do_restore() -> dict[str, Any]:
705
+ response = requests.post(
706
+ restore_url,
707
+ headers=headers,
708
+ json={"source_branch_id": branch.parent_id},
709
+ timeout=30,
710
+ )
711
+ response.raise_for_status()
712
+ return response.json()
713
+
714
+ # Wrap in retry logic to handle rate limits
715
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
716
+ data = _retry_on_rate_limit(do_restore, operation_name="branch_restore")
717
+ operations = data.get("operations", [])
718
+
719
+ # The restore API returns operations that run asynchronously.
720
+ # We must wait for operations to complete before the next test
721
+ # starts, otherwise connections may fail during the restore.
722
+ if operations:
723
+ _wait_for_operations(
724
+ project_id=branch.project_id,
725
+ operations=operations,
726
+ headers=headers,
727
+ base_url=base_url,
728
+ )
547
729
 
548
730
 
549
731
  def _wait_for_operations(
@@ -556,6 +738,9 @@ def _wait_for_operations(
556
738
  ) -> None:
557
739
  """Wait for Neon operations to complete.
558
740
 
741
+ Handles rate limit (429) errors with exponential backoff retry.
742
+ See: https://api-docs.neon.tech/reference/api-rate-limiting
743
+
559
744
  Args:
560
745
  project_id: The Neon project ID
561
746
  operations: List of operation dicts from the API response
@@ -589,10 +774,21 @@ def _wait_for_operations(
589
774
  still_pending = []
590
775
  for op_id in pending_op_ids:
591
776
  op_url = f"{base_url}/projects/{project_id}/operations/{op_id}"
592
- try:
593
- response = requests.get(op_url, headers=headers, timeout=10)
777
+
778
+ def get_operation_status(url: str = op_url) -> dict[str, Any]:
779
+ """Fetch operation status. Default arg captures url by value."""
780
+ response = requests.get(url, headers=headers, timeout=10)
594
781
  response.raise_for_status()
595
- op_data = response.json().get("operation", {})
782
+ return response.json()
783
+
784
+ try:
785
+ # Wrap in retry logic to handle rate limits
786
+ # See: https://api-docs.neon.tech/reference/api-rate-limiting
787
+ result = _retry_on_rate_limit(
788
+ get_operation_status,
789
+ operation_name=f"operation_status({op_id})",
790
+ )
791
+ op_data = result.get("operation", {})
596
792
  status = op_data.get("status")
597
793
 
598
794
  if status == "failed":
@@ -601,7 +797,7 @@ def _wait_for_operations(
601
797
  if status not in ("finished", "skipped", "cancelled"):
602
798
  still_pending.append(op_id)
603
799
  except requests.RequestException:
604
- # On network error, assume still pending and retry
800
+ # On network error (non-429), assume still pending and retry
605
801
  still_pending.append(op_id)
606
802
 
607
803
  pending_op_ids = still_pending
@@ -612,9 +808,24 @@ def _wait_for_operations(
612
808
  )
613
809
 
614
810
 
811
+ def _branch_to_dict(branch: NeonBranch) -> dict[str, Any]:
812
+ """Convert NeonBranch to a JSON-serializable dict."""
813
+ return asdict(branch)
814
+
815
+
816
+ def _dict_to_branch(data: dict[str, Any]) -> NeonBranch:
817
+ """Convert a dict back to NeonBranch."""
818
+ return NeonBranch(**data)
819
+
820
+
821
+ # Timeout for waiting for migrations to complete (seconds)
822
+ _MIGRATION_WAIT_TIMEOUT = 300 # 5 minutes
823
+
824
+
615
825
  @pytest.fixture(scope="session")
616
826
  def _neon_migration_branch(
617
827
  request: pytest.FixtureRequest,
828
+ tmp_path_factory: pytest.TempPathFactory,
618
829
  ) -> Generator[NeonBranch, None, None]:
619
830
  """
620
831
  Session-scoped branch where migrations are applied.
@@ -623,6 +834,13 @@ def _neon_migration_branch(
623
834
  the parent for all test branches. Migrations run once per session
624
835
  on this branch.
625
836
 
837
+ pytest-xdist Support:
838
+ When running with pytest-xdist, the first worker to acquire the lock
839
+ creates the migration branch. Other workers wait for migrations to
840
+ complete, then reuse the same branch. This avoids redundant API calls
841
+ and ensures migrations only run once. Only the creator cleans up the
842
+ branch at session end.
843
+
626
844
  Note: The migration branch cannot have an expiry because Neon doesn't
627
845
  allow creating child branches from branches with expiration dates.
628
846
  Cleanup relies on the fixture teardown at session end.
@@ -632,25 +850,102 @@ def _neon_migration_branch(
632
850
  it on request.config. After migrations run, _neon_branch_for_reset
633
851
  compares the fingerprint to detect if the schema actually changed.
634
852
  """
635
- # No expiry - Neon doesn't allow children from branches with expiry
636
- branch_gen = _create_neon_branch(
637
- request,
638
- branch_expiry_override=0,
639
- branch_name_suffix="-migrated",
853
+ config = request.config
854
+ worker_id = _get_xdist_worker_id()
855
+ is_xdist = worker_id != "main"
856
+
857
+ # Get env var name for DATABASE_URL
858
+ env_var_name = _get_config_value(
859
+ config, "neon_env_var", "", "neon_env_var", "DATABASE_URL"
640
860
  )
641
- branch = next(branch_gen)
642
861
 
643
- # Capture schema fingerprint BEFORE migrations run
644
- # This is stored on config so _neon_branch_for_reset can compare after
645
- pre_migration_fingerprint = _get_schema_fingerprint(branch.connection_string)
646
- request.config._neon_pre_migration_fingerprint = pre_migration_fingerprint # type: ignore[attr-defined]
862
+ # For xdist, use shared temp directory and filelock
863
+ # tmp_path_factory.getbasetemp().parent is shared across all workers
864
+ if is_xdist:
865
+ root_tmp_dir = tmp_path_factory.getbasetemp().parent
866
+ cache_file = root_tmp_dir / "neon_migration_branch.json"
867
+ lock_file = root_tmp_dir / "neon_migration_branch.lock"
868
+ else:
869
+ cache_file = None
870
+ lock_file = None
871
+
872
+ is_creator = False
873
+ branch: NeonBranch
874
+ branch_gen: Generator[NeonBranch, None, None] | None = None
875
+ original_env_value: str | None = None
876
+
877
+ if is_xdist:
878
+ assert cache_file is not None and lock_file is not None
879
+ with FileLock(str(lock_file)):
880
+ if cache_file.exists():
881
+ # Another worker already created the branch - reuse it
882
+ data = json.loads(cache_file.read_text())
883
+ branch = _dict_to_branch(data["branch"])
884
+ pre_migration_fingerprint = tuple(
885
+ tuple(row) for row in data["pre_migration_fingerprint"]
886
+ )
887
+ config._neon_pre_migration_fingerprint = pre_migration_fingerprint # type: ignore[attr-defined]
888
+
889
+ # Set DATABASE_URL for this worker (not done by _create_neon_branch)
890
+ original_env_value = os.environ.get(env_var_name)
891
+ os.environ[env_var_name] = branch.connection_string
892
+ else:
893
+ # First worker - create branch and cache it
894
+ is_creator = True
895
+ branch_gen = _create_neon_branch(
896
+ request,
897
+ branch_expiry_override=0,
898
+ branch_name_suffix="-migrated",
899
+ )
900
+ branch = next(branch_gen)
901
+
902
+ # Capture schema fingerprint BEFORE migrations run
903
+ pre_migration_fingerprint = _get_schema_fingerprint(
904
+ branch.connection_string
905
+ )
906
+ config._neon_pre_migration_fingerprint = pre_migration_fingerprint # type: ignore[attr-defined]
907
+
908
+ # Cache for other workers (they'll read this after lock released)
909
+ # Note: We cache now with pre-migration fingerprint. The branch
910
+ # content will have migrations applied by neon_apply_migrations.
911
+ cache_file.write_text(
912
+ json.dumps(
913
+ {
914
+ "branch": _branch_to_dict(branch),
915
+ "pre_migration_fingerprint": pre_migration_fingerprint,
916
+ }
917
+ )
918
+ )
919
+ else:
920
+ # Not using xdist - create branch normally
921
+ is_creator = True
922
+ branch_gen = _create_neon_branch(
923
+ request,
924
+ branch_expiry_override=0,
925
+ branch_name_suffix="-migrated",
926
+ )
927
+ branch = next(branch_gen)
928
+
929
+ # Capture schema fingerprint BEFORE migrations run
930
+ pre_migration_fingerprint = _get_schema_fingerprint(branch.connection_string)
931
+ config._neon_pre_migration_fingerprint = pre_migration_fingerprint # type: ignore[attr-defined]
932
+
933
+ # Mark whether this worker is the creator (used by neon_apply_migrations)
934
+ config._neon_is_migration_creator = is_creator # type: ignore[attr-defined]
647
935
 
648
936
  try:
649
937
  yield branch
650
938
  finally:
651
- # Clean up by exhausting the generator (triggers branch deletion)
652
- with contextlib.suppress(StopIteration):
653
- next(branch_gen)
939
+ # Restore env var if we set it (non-creator workers)
940
+ if original_env_value is not None:
941
+ os.environ[env_var_name] = original_env_value
942
+ elif not is_creator and env_var_name in os.environ:
943
+ os.environ.pop(env_var_name, None)
944
+
945
+ # Only the creator cleans up the branch
946
+ if is_creator and branch_gen is not None:
947
+ with contextlib.suppress(StopIteration):
948
+ next(branch_gen)
654
949
 
655
950
 
656
951
  @pytest.fixture(scope="session")
@@ -661,6 +956,12 @@ def neon_apply_migrations(_neon_migration_branch: NeonBranch) -> Any:
661
956
  The migration branch is already created and DATABASE_URL is set.
662
957
  Migrations run once per test session, before any tests execute.
663
958
 
959
+ pytest-xdist Support:
960
+ When running with pytest-xdist, migrations only run on the first
961
+ worker (the one that created the migration branch). Other workers
962
+ wait for migrations to complete before proceeding. This ensures
963
+ migrations run exactly once, even with parallel workers.
964
+
664
965
  Smart Migration Detection:
665
966
  The plugin automatically detects whether migrations actually modified
666
967
  the database schema. If no schema changes occurred (or this fixture
@@ -703,11 +1004,67 @@ def neon_apply_migrations(_neon_migration_branch: NeonBranch) -> Any:
703
1004
  return _MIGRATIONS_NOT_DEFINED
704
1005
 
705
1006
 
1007
+ @pytest.fixture(scope="session")
1008
+ def _neon_migrations_synchronized(
1009
+ request: pytest.FixtureRequest,
1010
+ tmp_path_factory: pytest.TempPathFactory,
1011
+ _neon_migration_branch: NeonBranch,
1012
+ neon_apply_migrations: Any,
1013
+ ) -> Any:
1014
+ """
1015
+ Internal fixture that synchronizes migrations across xdist workers.
1016
+
1017
+ This fixture ensures that:
1018
+ 1. Only the creator worker runs migrations
1019
+ 2. Other workers wait for migrations to complete before proceeding
1020
+ 3. The return value from neon_apply_migrations is preserved for detection
1021
+
1022
+ Without xdist, this is a simple passthrough.
1023
+ """
1024
+ config = request.config
1025
+ worker_id = _get_xdist_worker_id()
1026
+ is_xdist = worker_id != "main"
1027
+ is_creator = getattr(config, "_neon_is_migration_creator", True)
1028
+
1029
+ if not is_xdist:
1030
+ # Not using xdist - migrations already ran, just return the value
1031
+ return neon_apply_migrations
1032
+
1033
+ # For xdist, use a signal file to coordinate
1034
+ root_tmp_dir = tmp_path_factory.getbasetemp().parent
1035
+ migrations_done_file = root_tmp_dir / "neon_migrations_done"
1036
+ migrations_lock_file = root_tmp_dir / "neon_migrations.lock"
1037
+
1038
+ if is_creator:
1039
+ # Creator: migrations just ran via neon_apply_migrations dependency
1040
+ # Signal completion to other workers
1041
+ with FileLock(str(migrations_lock_file)):
1042
+ migrations_done_file.write_text("done")
1043
+ return neon_apply_migrations
1044
+ else:
1045
+ # Non-creator: wait for migrations to complete
1046
+ # The neon_apply_migrations fixture still runs but on already-migrated DB
1047
+ # (most migration tools handle this gracefully as a no-op)
1048
+ waited = 0.0
1049
+ poll_interval = 0.5
1050
+ while not migrations_done_file.exists():
1051
+ if waited >= _MIGRATION_WAIT_TIMEOUT:
1052
+ raise RuntimeError(
1053
+ f"Timeout waiting for migrations to complete after "
1054
+ f"{_MIGRATION_WAIT_TIMEOUT}s. The creator worker may have "
1055
+ f"failed or is still running migrations."
1056
+ )
1057
+ time.sleep(poll_interval)
1058
+ waited += poll_interval
1059
+
1060
+ return neon_apply_migrations
1061
+
1062
+
706
1063
  @pytest.fixture(scope="session")
707
1064
  def _neon_branch_for_reset(
708
1065
  request: pytest.FixtureRequest,
709
1066
  _neon_migration_branch: NeonBranch,
710
- neon_apply_migrations: Any, # Ensures migrations run first; value for detection
1067
+ _neon_migrations_synchronized: Any, # Ensures migrations complete; for detection
711
1068
  ) -> Generator[NeonBranch, None, None]:
712
1069
  """
713
1070
  Internal fixture that creates a test branch from the migration branch.
@@ -737,7 +1094,8 @@ def _neon_branch_for_reset(
737
1094
  - Migrations exist but are already applied (no schema changes)
738
1095
  """
739
1096
  # Check if migrations fixture was overridden
740
- migrations_defined = neon_apply_migrations is not _MIGRATIONS_NOT_DEFINED
1097
+ # _neon_migrations_synchronized passes through the neon_apply_migrations value
1098
+ migrations_defined = _neon_migrations_synchronized is not _MIGRATIONS_NOT_DEFINED
741
1099
 
742
1100
  # Check if schema actually changed (if we have a pre-migration fingerprint)
743
1101
  pre_fingerprint = getattr(request.config, "_neon_pre_migration_fingerprint", ())
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pytest-neon
3
- Version: 2.1.4
3
+ Version: 2.2.1
4
4
  Summary: Pytest plugin for Neon database branch isolation in tests
5
5
  Project-URL: Homepage, https://github.com/ZainRizvi/pytest-neon
6
6
  Project-URL: Repository, https://github.com/ZainRizvi/pytest-neon
@@ -24,6 +24,7 @@ Classifier: Programming Language :: Python :: 3.14
24
24
  Classifier: Topic :: Database
25
25
  Classifier: Topic :: Software Development :: Testing
26
26
  Requires-Python: >=3.9
27
+ Requires-Dist: filelock>=3.0
27
28
  Requires-Dist: neon-api>=0.1.0
28
29
  Requires-Dist: pytest>=7.0
29
30
  Requires-Dist: requests>=2.20
@@ -489,6 +490,21 @@ The `neon_branch_readwrite` fixture uses Neon's branch restore API to reset data
489
490
 
490
491
  This is similar to database transactions but at the branch level.
491
492
 
493
+ ## Branch Naming
494
+
495
+ Branches are automatically named to help identify their source:
496
+
497
+ ```
498
+ pytest-[git-branch]-[random]-[suffix]
499
+ ```
500
+
501
+ **Examples:**
502
+ - `pytest-main-a1b2-migrated` - Migration branch from `main`
503
+ - `pytest-feature-auth-c3d4-test-main` - Test branch from `feature/auth`
504
+ - `pytest-a1b2-migrated` - When not in a git repo
505
+
506
+ The git branch name is sanitized (only `a-z`, `0-9`, `-`, `_` allowed) and truncated to 15 characters. This makes it easy to identify orphaned branches in the Neon console.
507
+
492
508
  ## Parallel Test Execution (pytest-xdist)
493
509
 
494
510
  This plugin supports parallel test execution with [pytest-xdist](https://pytest-xdist.readthedocs.io/). Each xdist worker automatically gets its own isolated branch.
@@ -0,0 +1,8 @@
1
+ pytest_neon/__init__.py,sha256=MShSVrrcshqn0Gk9s9zc4iMopsVfqCbDyKYJhOa5QYE,398
2
+ pytest_neon/plugin.py,sha256=kpfRodNlIY_6a11UbvtF-S0KFjaPHPPo-K3kiaEAGsA,55952
3
+ pytest_neon/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
+ pytest_neon-2.2.1.dist-info/METADATA,sha256=NZEC6cCKDZ00UyyAjVQdyawyTjZ1MKL8oboOsw5ao0M,19266
5
+ pytest_neon-2.2.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
6
+ pytest_neon-2.2.1.dist-info/entry_points.txt,sha256=5U88Idj_G8-PSDb9VF3OwYFbGLHnGOo_GxgYvi0dtXw,37
7
+ pytest_neon-2.2.1.dist-info/licenses/LICENSE,sha256=aKKp_Ex4WBHTByY4BhXJ181dzB_qYhi2pCUmZ7Spn_0,1067
8
+ pytest_neon-2.2.1.dist-info/RECORD,,
@@ -1,8 +0,0 @@
1
- pytest_neon/__init__.py,sha256=bWWilGWaSJW2ofj0YsKqdC7r972s0c2Ar5DcE1Tn0Oc,398
2
- pytest_neon/plugin.py,sha256=A9qv1mv0CXez-SDAWruzVxjRFp0YdSWkU-8j9TWV_LA,41770
3
- pytest_neon/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
- pytest_neon-2.1.4.dist-info/METADATA,sha256=dmGCHSNi32pEA5wG6fZhjoFrcirza1kO16RdRSmyEJs,18734
5
- pytest_neon-2.1.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
6
- pytest_neon-2.1.4.dist-info/entry_points.txt,sha256=5U88Idj_G8-PSDb9VF3OwYFbGLHnGOo_GxgYvi0dtXw,37
7
- pytest_neon-2.1.4.dist-info/licenses/LICENSE,sha256=aKKp_Ex4WBHTByY4BhXJ181dzB_qYhi2pCUmZ7Spn_0,1067
8
- pytest_neon-2.1.4.dist-info/RECORD,,