pytest-neon 2.3.1__py3-none-any.whl → 3.0.0__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/plugin.py CHANGED
@@ -1,15 +1,10 @@
1
1
  """Pytest plugin providing Neon database branch fixtures.
2
2
 
3
- This plugin provides fixtures for database testing using Neon's instant
4
- branching feature. Multiple isolation levels are available:
3
+ This plugin provides a simple fixture for database testing using Neon's instant
4
+ branching feature. All tests share a single branch per session.
5
5
 
6
- Main fixtures:
7
- neon_branch_readonly: True read-only access via read_only endpoint (enforced)
8
- neon_branch_dirty: Session-scoped read-write, shared state across all tests
9
- neon_branch_isolated: Per-worker branch with reset after each test (recommended)
10
- neon_branch_readwrite: Deprecated, use neon_branch_isolated instead
11
- neon_branch: Deprecated alias for neon_branch_isolated
12
- neon_branch_shared: Shared branch without reset (module-scoped)
6
+ Main fixture:
7
+ neon_branch: Session-scoped shared branch for all tests
13
8
 
14
9
  Connection fixtures (require extras):
15
10
  neon_connection: psycopg2 connection (requires psycopg2 extra)
@@ -18,32 +13,35 @@ Connection fixtures (require extras):
18
13
 
19
14
  Architecture:
20
15
  Parent Branch (configured or project default)
21
- └── Migration Branch (session-scoped, read_write endpoint)
22
- ↑ migrations run here ONCE
23
-
24
- ├── Read-only Endpoint (read_only endpoint ON migration branch)
25
- │ ↑ neon_branch_readonly uses this
26
-
27
- ├── Dirty Branch (session-scoped child, shared across ALL workers)
28
- │ ↑ neon_branch_dirty uses this
29
-
30
- └── Isolated Branch (one per xdist worker, lazily created)
31
- ↑ neon_branch_isolated uses this, reset after each test
32
-
33
- SQLAlchemy Users:
34
- If you create your own SQLAlchemy engine (not using neon_engine fixture),
35
- you MUST use pool_pre_ping=True when using neon_branch_isolated:
36
-
37
- engine = create_engine(DATABASE_URL, pool_pre_ping=True)
38
-
39
- This is required because branch resets terminate server-side connections.
40
- Without pool_pre_ping, SQLAlchemy may try to reuse dead pooled connections,
41
- causing "SSL connection has been closed unexpectedly" errors.
16
+ └── Test Branch (session-scoped, 10-min expiry)
17
+ ↑ migrations run here ONCE, all tests share this
42
18
 
43
19
  Configuration:
44
20
  Set NEON_API_KEY and NEON_PROJECT_ID environment variables, or use
45
21
  --neon-api-key and --neon-project-id CLI options.
46
22
 
23
+ Test Isolation:
24
+ Since all tests share the same branch, tests that modify data will see
25
+ each other's changes. For test isolation, use one of these patterns:
26
+
27
+ 1. Transaction rollback:
28
+ @pytest.fixture
29
+ def db_transaction(neon_branch):
30
+ import psycopg
31
+ conn = psycopg.connect(neon_branch.connection_string)
32
+ conn.execute("BEGIN")
33
+ yield conn
34
+ conn.execute("ROLLBACK")
35
+ conn.close()
36
+
37
+ 2. Table truncation:
38
+ @pytest.fixture(autouse=True)
39
+ def clean_tables(neon_branch):
40
+ yield
41
+ with psycopg.connect(neon_branch.connection_string) as conn:
42
+ conn.execute("TRUNCATE users, orders CASCADE")
43
+ conn.commit()
44
+
47
45
  For full documentation, see: https://github.com/ZainRizvi/pytest-neon
48
46
  """
49
47
 
@@ -80,9 +78,6 @@ _RATE_LIMIT_MAX_TOTAL_DELAY = 90.0 # 1.5 minutes total cap
80
78
  _RATE_LIMIT_JITTER_FACTOR = 0.25 # +/- 25% jitter
81
79
  _RATE_LIMIT_MAX_ATTEMPTS = 10 # Maximum number of retry attempts
82
80
 
83
- # Sentinel value to detect when neon_apply_migrations was not overridden
84
- _MIGRATIONS_NOT_DEFINED = object()
85
-
86
81
 
87
82
  class NeonRateLimitError(Exception):
88
83
  """Raised when Neon API rate limit is exceeded and retries are exhausted."""
@@ -245,8 +240,8 @@ def _get_xdist_worker_id() -> str:
245
240
  Get the pytest-xdist worker ID, or "main" if not running under xdist.
246
241
 
247
242
  When running tests in parallel with pytest-xdist, each worker process
248
- gets a unique ID (gw0, gw1, gw2, etc.). This is used to create separate
249
- branches per worker to avoid database state pollution between parallel tests.
243
+ gets a unique ID (gw0, gw1, gw2, etc.). This is used to coordinate
244
+ branch creation and migrations across workers.
250
245
  """
251
246
  return os.environ.get("PYTEST_XDIST_WORKER", "main")
252
247
 
@@ -295,47 +290,35 @@ def _get_git_branch_name() -> str | None:
295
290
  return None
296
291
 
297
292
 
298
- def _extract_password_from_connection_string(connection_string: str) -> str:
299
- """Extract password from a PostgreSQL connection string."""
300
- # Format: postgresql://user:password@host/db?params
301
- from urllib.parse import urlparse
302
-
303
- parsed = urlparse(connection_string)
304
- if parsed.password:
305
- return parsed.password
306
- raise ValueError(f"No password found in connection string: {connection_string}")
307
-
308
-
309
- def _get_schema_fingerprint(connection_string: str) -> tuple[tuple[Any, ...], ...]:
293
+ def _reveal_role_password(
294
+ api_key: str, project_id: str, branch_id: str, role_name: str
295
+ ) -> str:
310
296
  """
311
- Get a fingerprint of the database schema for change detection.
297
+ Get the password for a role WITHOUT resetting it.
312
298
 
313
- Queries information_schema for all tables, columns, and their properties
314
- in the public schema. Returns a hashable tuple that can be compared
315
- before/after migrations to detect if the schema actually changed.
299
+ Uses Neon's reveal_password API endpoint (GET request).
316
300
 
317
- This is used to avoid creating unnecessary migration branches when
318
- no actual schema changes occurred.
301
+ Note: The neon-api library has a bug where it uses POST instead of GET,
302
+ so we make the request directly.
319
303
  """
304
+ url = (
305
+ f"https://console.neon.tech/api/v2/projects/{project_id}"
306
+ f"/branches/{branch_id}/roles/{role_name}/reveal_password"
307
+ )
308
+ headers = {
309
+ "Authorization": f"Bearer {api_key}",
310
+ "Accept": "application/json",
311
+ }
312
+
313
+ response = requests.get(url, headers=headers, timeout=30)
320
314
  try:
321
- import psycopg
322
- except ImportError:
323
- try:
324
- import psycopg2 as psycopg # type: ignore[import-not-found]
325
- except ImportError:
326
- # No driver available - can't fingerprint, assume migrations changed things
327
- return ()
328
-
329
- with psycopg.connect(connection_string) as conn, conn.cursor() as cur:
330
- cur.execute("""
331
- SELECT table_name, column_name, data_type, is_nullable,
332
- column_default, ordinal_position
333
- FROM information_schema.columns
334
- WHERE table_schema = 'public'
335
- ORDER BY table_name, ordinal_position
336
- """)
337
- rows = cur.fetchall()
338
- return tuple(tuple(row) for row in rows)
315
+ response.raise_for_status()
316
+ except requests.exceptions.HTTPError:
317
+ # Wrap in NeonAPIError for consistent error handling
318
+ raise NeonAPIError(response.text) from None
319
+
320
+ data = response.json()
321
+ return data["password"]
339
322
 
340
323
 
341
324
  @dataclass
@@ -448,7 +431,7 @@ class NeonBranchManager:
448
431
  Create a new Neon branch with a read_write endpoint.
449
432
 
450
433
  Args:
451
- name_suffix: Suffix to add to branch name (e.g., "-migration", "-dirty")
434
+ name_suffix: Suffix to add to branch name (e.g., "-test")
452
435
  parent_branch_id: Parent branch ID (defaults to config's parent)
453
436
  expiry_seconds: Branch expiry in seconds (0 or None for no expiry)
454
437
 
@@ -508,7 +491,7 @@ class NeonBranchManager:
508
491
  )
509
492
 
510
493
  # Get password
511
- connection_string = self._reset_password_and_build_connection_string(
494
+ connection_string = self._get_password_and_build_connection_string(
512
495
  branch.id, host
513
496
  )
514
497
 
@@ -521,53 +504,6 @@ class NeonBranchManager:
521
504
  endpoint_id=endpoint_id,
522
505
  )
523
506
 
524
- def create_readonly_endpoint(self, branch: NeonBranch) -> NeonBranch:
525
- """
526
- Create a read_only endpoint on an existing branch.
527
-
528
- This creates a true read-only endpoint that enforces no writes at the
529
- database level.
530
-
531
- Args:
532
- branch: The branch to create the endpoint on
533
-
534
- Returns:
535
- NeonBranch with the read_only endpoint's connection details
536
- """
537
- result = _retry_on_rate_limit(
538
- lambda: self._neon.endpoint_create(
539
- project_id=self.config.project_id,
540
- endpoint={
541
- "branch_id": branch.branch_id,
542
- "type": "read_only",
543
- },
544
- ),
545
- operation_name="endpoint_create_readonly",
546
- )
547
-
548
- endpoint_id = result.endpoint.id
549
- host = self._wait_for_endpoint(endpoint_id)
550
-
551
- # Reuse the password from the parent branch's connection string.
552
- # DO NOT call role_password_reset here - it would invalidate the
553
- # password used by the parent branch's read_write endpoint, breaking
554
- # any existing connections (especially in xdist where other workers
555
- # may be using the cached connection string).
556
- password = _extract_password_from_connection_string(branch.connection_string)
557
- connection_string = (
558
- f"postgresql://{self.config.role_name}:{password}@{host}/"
559
- f"{self.config.database_name}?sslmode=require"
560
- )
561
-
562
- return NeonBranch(
563
- branch_id=branch.branch_id,
564
- project_id=self.config.project_id,
565
- connection_string=connection_string,
566
- host=host,
567
- parent_id=branch.parent_id,
568
- endpoint_id=endpoint_id,
569
- )
570
-
571
507
  def delete_branch(self, branch_id: str) -> None:
572
508
  """Delete a branch (silently ignores errors)."""
573
509
  if self.config.keep_branches:
@@ -583,28 +519,6 @@ class NeonBranchManager:
583
519
  msg = f"Failed to delete Neon branch {branch_id}: {e}"
584
520
  warnings.warn(msg, stacklevel=2)
585
521
 
586
- def delete_endpoint(self, endpoint_id: str) -> None:
587
- """Delete an endpoint (silently ignores errors)."""
588
- try:
589
- _retry_on_rate_limit(
590
- lambda: self._neon.endpoint_delete(
591
- project_id=self.config.project_id, endpoint_id=endpoint_id
592
- ),
593
- operation_name="endpoint_delete",
594
- )
595
- except Exception as e:
596
- warnings.warn(
597
- f"Failed to delete Neon endpoint {endpoint_id}: {e}", stacklevel=2
598
- )
599
-
600
- def reset_branch(self, branch: NeonBranch) -> None:
601
- """Reset a branch to its parent's state."""
602
- if not branch.parent_id:
603
- msg = f"Branch {branch.branch_id} has no parent - cannot reset"
604
- raise RuntimeError(msg)
605
-
606
- _reset_branch_to_parent(branch, self.config.api_key)
607
-
608
522
  def _wait_for_endpoint(self, endpoint_id: str, max_wait_seconds: float = 60) -> str:
609
523
  """Wait for endpoint to become active and return its host."""
610
524
  poll_interval = 0.5
@@ -632,19 +546,19 @@ class NeonBranchManager:
632
546
  time.sleep(poll_interval)
633
547
  waited += poll_interval
634
548
 
635
- def _reset_password_and_build_connection_string(
549
+ def _get_password_and_build_connection_string(
636
550
  self, branch_id: str, host: str
637
551
  ) -> str:
638
- """Reset role password and build connection string."""
639
- password_response = _retry_on_rate_limit(
640
- lambda: self._neon.role_password_reset(
552
+ """Get role password (without resetting) and build connection string."""
553
+ password = _retry_on_rate_limit(
554
+ lambda: _reveal_role_password(
555
+ api_key=self.config.api_key,
641
556
  project_id=self.config.project_id,
642
557
  branch_id=branch_id,
643
558
  role_name=self.config.role_name,
644
559
  ),
645
- operation_name="role_password_reset",
560
+ operation_name="role_password_reveal",
646
561
  )
647
- password = password_response.role.password
648
562
 
649
563
  return (
650
564
  f"postgresql://{self.config.role_name}:{password}@{host}/"
@@ -657,7 +571,7 @@ class XdistCoordinator:
657
571
  Coordinates branch sharing across pytest-xdist workers.
658
572
 
659
573
  Uses file locks and JSON cache files to ensure only one worker creates
660
- shared resources (like the migration branch), while others reuse them.
574
+ shared resources (like the test branch), while others reuse them.
661
575
  """
662
576
 
663
577
  def __init__(self, tmp_path_factory: pytest.TempPathFactory):
@@ -771,8 +685,7 @@ def _get_default_branch_id(neon: NeonAPI, project_id: str) -> str | None:
771
685
  Get the default/primary branch ID for a project.
772
686
 
773
687
  This is used as a safety check to ensure we never accidentally
774
- perform destructive operations (like password reset) on the
775
- production branch.
688
+ perform destructive operations on the production branch.
776
689
 
777
690
  Returns:
778
691
  The branch ID of the default branch, or None if not found.
@@ -901,461 +814,6 @@ def _get_config_value(
901
814
  return default
902
815
 
903
816
 
904
- def _create_neon_branch(
905
- request: pytest.FixtureRequest,
906
- parent_branch_id_override: str | None = None,
907
- branch_expiry_override: int | None = None,
908
- branch_name_suffix: str = "",
909
- ) -> Generator[NeonBranch, None, None]:
910
- """
911
- Internal helper that creates and manages a Neon branch lifecycle.
912
-
913
- This is the core implementation used by branch fixtures.
914
-
915
- Args:
916
- request: Pytest fixture request
917
- parent_branch_id_override: If provided, use this as parent instead of config
918
- branch_expiry_override: If provided, use this expiry instead of config
919
- branch_name_suffix: Optional suffix for branch name (e.g., "-migrated", "-test")
920
- """
921
- config = request.config
922
-
923
- api_key = _get_config_value(config, "neon_api_key", "NEON_API_KEY", "neon_api_key")
924
- project_id = _get_config_value(
925
- config, "neon_project_id", "NEON_PROJECT_ID", "neon_project_id"
926
- )
927
- # Use override if provided, otherwise read from config
928
- parent_branch_id = parent_branch_id_override or _get_config_value(
929
- config, "neon_parent_branch", "NEON_PARENT_BRANCH_ID", "neon_parent_branch"
930
- )
931
- database_name = _get_config_value(
932
- config, "neon_database", "NEON_DATABASE", "neon_database", "neondb"
933
- )
934
- role_name = _get_config_value(
935
- config, "neon_role", "NEON_ROLE", "neon_role", "neondb_owner"
936
- )
937
-
938
- # For boolean/int options, check CLI first, then ini
939
- keep_branches = config.getoption("neon_keep_branches", default=None)
940
- if keep_branches is None:
941
- keep_branches = config.getini("neon_keep_branches")
942
-
943
- # Use override if provided, otherwise read from config
944
- if branch_expiry_override is not None:
945
- branch_expiry = branch_expiry_override
946
- else:
947
- branch_expiry = config.getoption("neon_branch_expiry", default=None)
948
- if branch_expiry is None:
949
- branch_expiry = int(config.getini("neon_branch_expiry"))
950
-
951
- env_var_name = _get_config_value(
952
- config, "neon_env_var", "", "neon_env_var", "DATABASE_URL"
953
- )
954
-
955
- if not api_key:
956
- pytest.skip(
957
- "Neon API key not configured (set NEON_API_KEY or use --neon-api-key)"
958
- )
959
- if not project_id:
960
- pytest.skip(
961
- "Neon project ID not configured "
962
- "(set NEON_PROJECT_ID or use --neon-project-id)"
963
- )
964
-
965
- neon = NeonAPI(api_key=api_key)
966
-
967
- # Cache the default branch ID for safety checks (only fetch once per session)
968
- if not hasattr(config, "_neon_default_branch_id"):
969
- config._neon_default_branch_id = _get_default_branch_id(neon, project_id) # type: ignore[attr-defined]
970
-
971
- # Generate unique branch name
972
- # Format: pytest-[git branch (first 15 chars)]-[random]-[suffix]
973
- # This helps identify orphaned branches by showing which git branch created them
974
- random_suffix = os.urandom(2).hex() # 2 bytes = 4 hex chars
975
- git_branch = _get_git_branch_name()
976
- if git_branch:
977
- # Truncate git branch to 15 chars to keep branch names reasonable
978
- git_prefix = git_branch[:15]
979
- branch_name = f"pytest-{git_prefix}-{random_suffix}{branch_name_suffix}"
980
- else:
981
- branch_name = f"pytest-{random_suffix}{branch_name_suffix}"
982
-
983
- # Build branch creation payload
984
- branch_config: dict[str, Any] = {"name": branch_name}
985
- if parent_branch_id:
986
- branch_config["parent_id"] = parent_branch_id
987
-
988
- # Set branch expiration (auto-delete) as a safety net for interrupted test runs
989
- # This uses the branch expires_at field, not endpoint suspend_timeout
990
- if branch_expiry and branch_expiry > 0:
991
- expires_at = datetime.now(timezone.utc) + timedelta(seconds=branch_expiry)
992
- branch_config["expires_at"] = expires_at.strftime("%Y-%m-%dT%H:%M:%SZ")
993
-
994
- # Create branch with compute endpoint
995
- # Wrap in retry logic to handle rate limits
996
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
997
- result = _retry_on_rate_limit(
998
- lambda: neon.branch_create(
999
- project_id=project_id,
1000
- branch=branch_config,
1001
- endpoints=[{"type": "read_write"}],
1002
- ),
1003
- operation_name="branch_create",
1004
- )
1005
-
1006
- branch = result.branch
1007
-
1008
- # Get endpoint_id from operations
1009
- # (branch_create returns operations, not endpoints directly)
1010
- endpoint_id = None
1011
- for op in result.operations:
1012
- if op.endpoint_id:
1013
- endpoint_id = op.endpoint_id
1014
- break
1015
-
1016
- if not endpoint_id:
1017
- raise RuntimeError(f"No endpoint created for branch {branch.id}")
1018
-
1019
- # Wait for endpoint to be ready (it starts in "init" state)
1020
- # Endpoints typically become active in 1-2 seconds, but we allow up to 60s
1021
- # to handle occasional Neon API slowness or high load scenarios
1022
- max_wait_seconds = 60
1023
- poll_interval = 0.5 # Poll every 500ms for responsive feedback
1024
- waited = 0.0
1025
-
1026
- while True:
1027
- # Wrap in retry logic to handle rate limits during polling
1028
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1029
- endpoint_response = _retry_on_rate_limit(
1030
- lambda: neon.endpoint(project_id=project_id, endpoint_id=endpoint_id),
1031
- operation_name="endpoint_status",
1032
- )
1033
- endpoint = endpoint_response.endpoint
1034
- state = endpoint.current_state
1035
-
1036
- if state == EndpointState.active:
1037
- break
1038
-
1039
- if waited >= max_wait_seconds:
1040
- raise RuntimeError(
1041
- f"Timeout waiting for endpoint {endpoint_id} to become active "
1042
- f"(current state: {state})"
1043
- )
1044
-
1045
- time.sleep(poll_interval)
1046
- waited += poll_interval
1047
-
1048
- host = endpoint.host
1049
-
1050
- # SAFETY CHECK: Ensure we never reset password on the default/production branch
1051
- # This should be impossible since we just created this branch, but we check
1052
- # defensively to prevent catastrophic mistakes if there's ever a bug
1053
- default_branch_id = getattr(config, "_neon_default_branch_id", None)
1054
- if default_branch_id and branch.id == default_branch_id:
1055
- raise RuntimeError(
1056
- f"SAFETY CHECK FAILED: Attempted to reset password on default branch "
1057
- f"{branch.id}. This should never happen - the plugin creates new "
1058
- f"branches and should never operate on the default branch. "
1059
- f"Please report this bug at https://github.com/ZainRizvi/pytest-neon/issues"
1060
- )
1061
-
1062
- # Reset password to get the password value
1063
- # (newly created branches don't expose password)
1064
- # Wrap in retry logic to handle rate limits
1065
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1066
- password_response = _retry_on_rate_limit(
1067
- lambda: neon.role_password_reset(
1068
- project_id=project_id,
1069
- branch_id=branch.id,
1070
- role_name=role_name,
1071
- ),
1072
- operation_name="role_password_reset",
1073
- )
1074
- password = password_response.role.password
1075
-
1076
- # Build connection string
1077
- connection_string = (
1078
- f"postgresql://{role_name}:{password}@{host}/{database_name}?sslmode=require"
1079
- )
1080
-
1081
- neon_branch_info = NeonBranch(
1082
- branch_id=branch.id,
1083
- project_id=project_id,
1084
- connection_string=connection_string,
1085
- host=host,
1086
- parent_id=branch.parent_id,
1087
- endpoint_id=endpoint_id,
1088
- )
1089
-
1090
- # Set DATABASE_URL (or configured env var) for the duration of the fixture scope
1091
- original_env_value = os.environ.get(env_var_name)
1092
- os.environ[env_var_name] = connection_string
1093
-
1094
- try:
1095
- yield neon_branch_info
1096
- finally:
1097
- # Restore original env var
1098
- if original_env_value is None:
1099
- os.environ.pop(env_var_name, None)
1100
- else:
1101
- os.environ[env_var_name] = original_env_value
1102
-
1103
- # Cleanup: delete branch unless --neon-keep-branches was specified
1104
- if not keep_branches:
1105
- try:
1106
- # Wrap in retry logic to handle rate limits
1107
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1108
- _retry_on_rate_limit(
1109
- lambda: neon.branch_delete(
1110
- project_id=project_id, branch_id=branch.id
1111
- ),
1112
- operation_name="branch_delete",
1113
- )
1114
- except Exception as e:
1115
- # Log but don't fail tests due to cleanup issues
1116
- warnings.warn(
1117
- f"Failed to delete Neon branch {branch.id}: {e}",
1118
- stacklevel=2,
1119
- )
1120
-
1121
-
1122
- def _create_readonly_endpoint(
1123
- branch: NeonBranch,
1124
- api_key: str,
1125
- database_name: str,
1126
- role_name: str,
1127
- ) -> NeonBranch:
1128
- """
1129
- Create a read_only endpoint on an existing branch.
1130
-
1131
- Returns a new NeonBranch object with the read_only endpoint's connection string.
1132
- The read_only endpoint enforces that no writes can be made through this connection.
1133
-
1134
- Args:
1135
- branch: The branch to create a read_only endpoint on
1136
- api_key: Neon API key
1137
- database_name: Database name for connection string
1138
- role_name: Role name for connection string
1139
-
1140
- Returns:
1141
- NeonBranch with read_only endpoint connection details
1142
- """
1143
- neon = NeonAPI(api_key=api_key)
1144
-
1145
- # Create read_only endpoint on the branch
1146
- # See: https://api-docs.neon.tech/reference/createprojectendpoint
1147
- result = _retry_on_rate_limit(
1148
- lambda: neon.endpoint_create(
1149
- project_id=branch.project_id,
1150
- endpoint={
1151
- "branch_id": branch.branch_id,
1152
- "type": "read_only",
1153
- },
1154
- ),
1155
- operation_name="endpoint_create_readonly",
1156
- )
1157
-
1158
- endpoint = result.endpoint
1159
- endpoint_id = endpoint.id
1160
-
1161
- # Wait for endpoint to be ready
1162
- max_wait_seconds = 60
1163
- poll_interval = 0.5
1164
- waited = 0.0
1165
-
1166
- while True:
1167
- endpoint_response = _retry_on_rate_limit(
1168
- lambda: neon.endpoint(
1169
- project_id=branch.project_id, endpoint_id=endpoint_id
1170
- ),
1171
- operation_name="endpoint_status_readonly",
1172
- )
1173
- endpoint = endpoint_response.endpoint
1174
- state = endpoint.current_state
1175
-
1176
- if state == EndpointState.active:
1177
- break
1178
-
1179
- if waited >= max_wait_seconds:
1180
- raise RuntimeError(
1181
- f"Timeout waiting for read_only endpoint {endpoint_id} "
1182
- f"to become active (current state: {state})"
1183
- )
1184
-
1185
- time.sleep(poll_interval)
1186
- waited += poll_interval
1187
-
1188
- host = endpoint.host
1189
-
1190
- # Reuse the password from the parent branch's connection string.
1191
- # DO NOT call role_password_reset here - it would invalidate the
1192
- # password used by the parent branch's read_write endpoint.
1193
- password = _extract_password_from_connection_string(branch.connection_string)
1194
-
1195
- # Build connection string for the read_only endpoint
1196
- connection_string = (
1197
- f"postgresql://{role_name}:{password}@{host}/{database_name}?sslmode=require"
1198
- )
1199
-
1200
- return NeonBranch(
1201
- branch_id=branch.branch_id,
1202
- project_id=branch.project_id,
1203
- connection_string=connection_string,
1204
- host=host,
1205
- parent_id=branch.parent_id,
1206
- endpoint_id=endpoint_id,
1207
- )
1208
-
1209
-
1210
- def _delete_endpoint(project_id: str, endpoint_id: str, api_key: str) -> None:
1211
- """Delete a Neon endpoint."""
1212
- neon = NeonAPI(api_key=api_key)
1213
- try:
1214
- _retry_on_rate_limit(
1215
- lambda: neon.endpoint_delete(
1216
- project_id=project_id, endpoint_id=endpoint_id
1217
- ),
1218
- operation_name="endpoint_delete",
1219
- )
1220
- except Exception as e:
1221
- warnings.warn(
1222
- f"Failed to delete Neon endpoint {endpoint_id}: {e}",
1223
- stacklevel=2,
1224
- )
1225
-
1226
-
1227
- def _reset_branch_to_parent(branch: NeonBranch, api_key: str) -> None:
1228
- """Reset a branch to its parent's state using the Neon API.
1229
-
1230
- Uses exponential backoff retry logic with jitter to handle rate limit (429)
1231
- errors. After initiating the restore, polls the operation status until it
1232
- completes.
1233
-
1234
- See: https://api-docs.neon.tech/reference/api-rate-limiting
1235
-
1236
- Args:
1237
- branch: The branch to reset
1238
- api_key: Neon API key
1239
- """
1240
- if not branch.parent_id:
1241
- raise RuntimeError(f"Branch {branch.branch_id} has no parent - cannot reset")
1242
-
1243
- base_url = "https://console.neon.tech/api/v2"
1244
- project_id = branch.project_id
1245
- branch_id = branch.branch_id
1246
- restore_url = f"{base_url}/projects/{project_id}/branches/{branch_id}/restore"
1247
- headers = {
1248
- "Authorization": f"Bearer {api_key}",
1249
- "Content-Type": "application/json",
1250
- }
1251
-
1252
- def do_restore() -> dict[str, Any]:
1253
- response = requests.post(
1254
- restore_url,
1255
- headers=headers,
1256
- json={"source_branch_id": branch.parent_id},
1257
- timeout=30,
1258
- )
1259
- response.raise_for_status()
1260
- return response.json()
1261
-
1262
- # Wrap in retry logic to handle rate limits
1263
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1264
- data = _retry_on_rate_limit(do_restore, operation_name="branch_restore")
1265
- operations = data.get("operations", [])
1266
-
1267
- # The restore API returns operations that run asynchronously.
1268
- # We must wait for operations to complete before the next test
1269
- # starts, otherwise connections may fail during the restore.
1270
- if operations:
1271
- _wait_for_operations(
1272
- project_id=branch.project_id,
1273
- operations=operations,
1274
- headers=headers,
1275
- base_url=base_url,
1276
- )
1277
-
1278
-
1279
- def _wait_for_operations(
1280
- project_id: str,
1281
- operations: list[dict[str, Any]],
1282
- headers: dict[str, str],
1283
- base_url: str,
1284
- max_wait_seconds: float = 60,
1285
- poll_interval: float = 0.5,
1286
- ) -> None:
1287
- """Wait for Neon operations to complete.
1288
-
1289
- Handles rate limit (429) errors with exponential backoff retry.
1290
- See: https://api-docs.neon.tech/reference/api-rate-limiting
1291
-
1292
- Args:
1293
- project_id: The Neon project ID
1294
- operations: List of operation dicts from the API response
1295
- headers: HTTP headers including auth
1296
- base_url: Base URL for Neon API
1297
- max_wait_seconds: Maximum time to wait (default: 60s)
1298
- poll_interval: Time between polls (default: 0.5s)
1299
- """
1300
- # Get operation IDs that aren't already finished
1301
- pending_op_ids = [
1302
- op["id"] for op in operations if op.get("status") not in ("finished", "skipped")
1303
- ]
1304
-
1305
- if not pending_op_ids:
1306
- return # All operations already complete
1307
-
1308
- waited = 0.0
1309
- first_poll = True
1310
- while pending_op_ids and waited < max_wait_seconds:
1311
- # Poll immediately first time (operation usually completes instantly),
1312
- # then wait between subsequent polls
1313
- if first_poll:
1314
- time.sleep(0.1) # Tiny delay to let operation start
1315
- waited += 0.1
1316
- first_poll = False
1317
- else:
1318
- time.sleep(poll_interval)
1319
- waited += poll_interval
1320
-
1321
- # Check status of each pending operation
1322
- still_pending = []
1323
- for op_id in pending_op_ids:
1324
- op_url = f"{base_url}/projects/{project_id}/operations/{op_id}"
1325
-
1326
- def get_operation_status(url: str = op_url) -> dict[str, Any]:
1327
- """Fetch operation status. Default arg captures url by value."""
1328
- response = requests.get(url, headers=headers, timeout=10)
1329
- response.raise_for_status()
1330
- return response.json()
1331
-
1332
- try:
1333
- # Wrap in retry logic to handle rate limits
1334
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1335
- result = _retry_on_rate_limit(
1336
- get_operation_status,
1337
- operation_name=f"operation_status({op_id})",
1338
- )
1339
- op_data = result.get("operation", {})
1340
- status = op_data.get("status")
1341
-
1342
- if status == "failed":
1343
- err = op_data.get("error", "unknown error")
1344
- raise RuntimeError(f"Operation {op_id} failed: {err}")
1345
- if status not in ("finished", "skipped", "cancelled"):
1346
- still_pending.append(op_id)
1347
- except requests.RequestException:
1348
- # On network error (non-429), assume still pending and retry
1349
- still_pending.append(op_id)
1350
-
1351
- pending_op_ids = still_pending
1352
-
1353
- if pending_op_ids:
1354
- raise RuntimeError(
1355
- f"Timeout waiting for operations to complete: {pending_op_ids}"
1356
- )
1357
-
1358
-
1359
817
  def _branch_to_dict(branch: NeonBranch) -> dict[str, Any]:
1360
818
  """Convert NeonBranch to a JSON-serializable dict."""
1361
819
  return asdict(branch)
@@ -1401,301 +859,110 @@ def _neon_xdist_coordinator(
1401
859
 
1402
860
 
1403
861
  @pytest.fixture(scope="session")
1404
- def _neon_migration_branch(
1405
- request: pytest.FixtureRequest,
862
+ def _neon_test_branch(
1406
863
  _neon_config: NeonConfig,
1407
864
  _neon_branch_manager: NeonBranchManager,
1408
865
  _neon_xdist_coordinator: XdistCoordinator,
1409
- ) -> Generator[NeonBranch, None, None]:
866
+ ) -> Generator[tuple[NeonBranch, bool], None, None]:
1410
867
  """
1411
- Session-scoped branch where migrations are applied.
868
+ Internal: Create test branch, coordinated across workers.
1412
869
 
1413
- This branch is ALWAYS created from the configured parent and serves as
1414
- the parent for all test branches (dirty, isolated, readonly endpoint).
1415
- Migrations run once per session on this branch.
870
+ This creates a single branch with expiry that all tests share.
871
+ The first worker creates the branch, others reuse it.
1416
872
 
1417
- pytest-xdist Support:
1418
- When running with pytest-xdist, the first worker to acquire the lock
1419
- creates the migration branch. Other workers wait for migrations to
1420
- complete, then reuse the same branch. This avoids redundant API calls
1421
- and ensures migrations only run once. Only the creator cleans up the
1422
- branch at session end.
1423
-
1424
- Note: The migration branch cannot have an expiry because Neon doesn't
1425
- allow creating child branches from branches with expiration dates.
1426
- Cleanup relies on the fixture teardown at session end.
873
+ Yields:
874
+ Tuple of (branch, is_creator) where is_creator indicates if this
875
+ worker created the branch (and should run migrations/cleanup).
1427
876
  """
1428
877
  env_manager = EnvironmentManager(_neon_config.env_var_name)
1429
- branch: NeonBranch
1430
- is_creator: bool
1431
878
 
1432
- def create_migration_branch() -> dict[str, Any]:
879
+ def create_branch() -> dict[str, Any]:
1433
880
  b = _neon_branch_manager.create_branch(
1434
- name_suffix="-migration",
1435
- expiry_seconds=0, # No expiry - child branches need this
881
+ name_suffix="-test",
882
+ expiry_seconds=_neon_config.branch_expiry,
1436
883
  )
1437
884
  return {"branch": _branch_to_dict(b)}
1438
885
 
1439
- # Coordinate branch creation across xdist workers
1440
886
  data, is_creator = _neon_xdist_coordinator.coordinate_resource(
1441
- "migration_branch", create_migration_branch
887
+ "test_branch", create_branch
1442
888
  )
1443
889
  branch = _dict_to_branch(data["branch"])
1444
-
1445
- # Store creator status for other fixtures
1446
- request.config._neon_is_migration_creator = is_creator # type: ignore[attr-defined]
1447
-
1448
- # Set DATABASE_URL
1449
890
  env_manager.set(branch.connection_string)
1450
891
 
1451
- # Non-creators wait for migrations to complete
1452
- if not is_creator:
1453
- _neon_xdist_coordinator.wait_for_signal(
1454
- "migrations_done", timeout=_MIGRATION_WAIT_TIMEOUT
1455
- )
1456
-
1457
892
  try:
1458
- yield branch
893
+ yield branch, is_creator
1459
894
  finally:
1460
895
  env_manager.restore()
1461
- # Only creator cleans up
1462
896
  if is_creator:
1463
897
  _neon_branch_manager.delete_branch(branch.branch_id)
1464
898
 
1465
899
 
1466
900
  @pytest.fixture(scope="session")
1467
- def neon_apply_migrations(_neon_migration_branch: NeonBranch) -> Any:
901
+ def neon_apply_migrations(_neon_test_branch: tuple[NeonBranch, bool]) -> Any:
1468
902
  """
1469
903
  Override this fixture to run migrations on the test database.
1470
904
 
1471
- The migration branch is already created and DATABASE_URL is set.
905
+ The test branch is already created and DATABASE_URL is set.
1472
906
  Migrations run once per test session, before any tests execute.
1473
907
 
1474
908
  pytest-xdist Support:
1475
909
  When running with pytest-xdist, migrations only run on the first
1476
- worker (the one that created the migration branch). Other workers
910
+ worker (the one that created the test branch). Other workers
1477
911
  wait for migrations to complete before proceeding. This ensures
1478
912
  migrations run exactly once, even with parallel workers.
1479
913
 
1480
- Smart Migration Detection:
1481
- The plugin automatically detects whether migrations actually modified
1482
- the database schema. If no schema changes occurred (or this fixture
1483
- isn't overridden), the plugin skips creating a separate migration
1484
- branch, saving Neon costs and branch slots.
1485
-
1486
914
  Example in conftest.py:
1487
915
 
1488
916
  @pytest.fixture(scope="session")
1489
- def neon_apply_migrations(_neon_migration_branch):
917
+ def neon_apply_migrations(_neon_test_branch):
1490
918
  import subprocess
1491
919
  subprocess.run(["alembic", "upgrade", "head"], check=True)
1492
920
 
1493
921
  Or with Django:
1494
922
 
1495
923
  @pytest.fixture(scope="session")
1496
- def neon_apply_migrations(_neon_migration_branch):
924
+ def neon_apply_migrations(_neon_test_branch):
1497
925
  from django.core.management import call_command
1498
926
  call_command("migrate", "--noinput")
1499
927
 
1500
928
  Or with raw SQL:
1501
929
 
1502
930
  @pytest.fixture(scope="session")
1503
- def neon_apply_migrations(_neon_migration_branch):
931
+ def neon_apply_migrations(_neon_test_branch):
1504
932
  import psycopg
1505
- with psycopg.connect(_neon_migration_branch.connection_string) as conn:
933
+ branch, is_creator = _neon_test_branch
934
+ with psycopg.connect(branch.connection_string) as conn:
1506
935
  with open("schema.sql") as f:
1507
936
  conn.execute(f.read())
1508
937
  conn.commit()
1509
938
 
1510
939
  Args:
1511
- _neon_migration_branch: The migration branch with connection details.
1512
- Use _neon_migration_branch.connection_string to connect directly,
940
+ _neon_test_branch: Tuple of (NeonBranch, is_creator).
941
+ Use _neon_test_branch[0].connection_string to connect directly,
1513
942
  or rely on DATABASE_URL which is already set.
1514
943
 
1515
944
  Returns:
1516
- Any value (ignored). The default returns a sentinel to indicate
1517
- the fixture was not overridden.
945
+ Any value (ignored). The default returns None.
1518
946
  """
1519
- return _MIGRATIONS_NOT_DEFINED
947
+ return None
1520
948
 
1521
949
 
1522
950
  @pytest.fixture(scope="session")
1523
- def _neon_migrations_synchronized(
1524
- request: pytest.FixtureRequest,
1525
- _neon_migration_branch: NeonBranch,
951
+ def neon_branch(
952
+ _neon_test_branch: tuple[NeonBranch, bool],
1526
953
  _neon_xdist_coordinator: XdistCoordinator,
1527
954
  neon_apply_migrations: Any,
1528
- ) -> Any:
1529
- """
1530
- Internal fixture that synchronizes migrations across xdist workers.
1531
-
1532
- This fixture ensures that:
1533
- 1. Only the creator worker runs migrations (non-creators wait in
1534
- _neon_migration_branch BEFORE neon_apply_migrations runs)
1535
- 2. Creator signals completion after migrations finish
1536
- 3. The return value from neon_apply_migrations is preserved for detection
1537
-
1538
- Without xdist, this is a simple passthrough.
1539
- """
1540
- is_creator = getattr(request.config, "_neon_is_migration_creator", True)
1541
-
1542
- if is_creator:
1543
- # Creator: migrations just ran via neon_apply_migrations dependency
1544
- # Signal completion to other workers
1545
- _neon_xdist_coordinator.send_signal("migrations_done")
1546
-
1547
- return neon_apply_migrations
1548
-
1549
-
1550
- @pytest.fixture(scope="session")
1551
- def _neon_dirty_branch(
1552
- _neon_config: NeonConfig,
1553
- _neon_branch_manager: NeonBranchManager,
1554
- _neon_xdist_coordinator: XdistCoordinator,
1555
- _neon_migration_branch: NeonBranch,
1556
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1557
- ) -> Generator[NeonBranch, None, None]:
1558
- """
1559
- Session-scoped dirty branch shared across ALL xdist workers.
1560
-
1561
- This branch is a child of the migration branch. All tests using
1562
- neon_branch_dirty share this single branch - writes persist and
1563
- are visible to all tests (even across workers).
1564
-
1565
- This is the "dirty" branch because:
1566
- - No reset between tests
1567
- - Shared across all workers (concurrent writes possible)
1568
- - Fast because no per-test overhead
1569
- """
1570
- env_manager = EnvironmentManager(_neon_config.env_var_name)
1571
- branch: NeonBranch
1572
- is_creator: bool
1573
-
1574
- def create_dirty_branch() -> dict[str, Any]:
1575
- b = _neon_branch_manager.create_branch(
1576
- name_suffix="-dirty",
1577
- parent_branch_id=_neon_migration_branch.branch_id,
1578
- expiry_seconds=_neon_config.branch_expiry,
1579
- )
1580
- return {"branch": _branch_to_dict(b)}
1581
-
1582
- # Coordinate dirty branch creation - shared across ALL workers
1583
- data, is_creator = _neon_xdist_coordinator.coordinate_resource(
1584
- "dirty_branch", create_dirty_branch
1585
- )
1586
- branch = _dict_to_branch(data["branch"])
1587
-
1588
- # Set DATABASE_URL
1589
- env_manager.set(branch.connection_string)
1590
-
1591
- try:
1592
- yield branch
1593
- finally:
1594
- env_manager.restore()
1595
- # Only creator cleans up
1596
- if is_creator:
1597
- _neon_branch_manager.delete_branch(branch.branch_id)
1598
-
1599
-
1600
- @pytest.fixture(scope="session")
1601
- def _neon_readonly_endpoint(
1602
- _neon_config: NeonConfig,
1603
- _neon_branch_manager: NeonBranchManager,
1604
- _neon_xdist_coordinator: XdistCoordinator,
1605
- _neon_migration_branch: NeonBranch,
1606
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1607
- ) -> Generator[NeonBranch, None, None]:
1608
- """
1609
- Session-scoped read_only endpoint on the migration branch.
1610
-
1611
- This is a true read-only endpoint - writes are blocked at the database
1612
- level. All workers share this endpoint since it's read-only anyway.
1613
- """
1614
- env_manager = EnvironmentManager(_neon_config.env_var_name)
1615
- branch: NeonBranch
1616
- is_creator: bool
1617
-
1618
- def create_readonly_endpoint() -> dict[str, Any]:
1619
- b = _neon_branch_manager.create_readonly_endpoint(_neon_migration_branch)
1620
- return {"branch": _branch_to_dict(b)}
1621
-
1622
- # Coordinate endpoint creation - shared across ALL workers
1623
- data, is_creator = _neon_xdist_coordinator.coordinate_resource(
1624
- "readonly_endpoint", create_readonly_endpoint
1625
- )
1626
- branch = _dict_to_branch(data["branch"])
1627
-
1628
- # Set DATABASE_URL
1629
- env_manager.set(branch.connection_string)
1630
-
1631
- try:
1632
- yield branch
1633
- finally:
1634
- env_manager.restore()
1635
- # Only creator cleans up the endpoint
1636
- if is_creator and branch.endpoint_id:
1637
- _neon_branch_manager.delete_endpoint(branch.endpoint_id)
1638
-
1639
-
1640
- @pytest.fixture(scope="session")
1641
- def _neon_isolated_branch(
1642
- request: pytest.FixtureRequest,
1643
- _neon_config: NeonConfig,
1644
- _neon_branch_manager: NeonBranchManager,
1645
- _neon_xdist_coordinator: XdistCoordinator,
1646
- _neon_migration_branch: NeonBranch,
1647
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1648
- ) -> Generator[NeonBranch, None, None]:
1649
- """
1650
- Session-scoped isolated branch, one per xdist worker.
1651
-
1652
- Each worker gets its own branch. Unlike the dirty branch, this is
1653
- per-worker to allow reset operations without affecting other workers.
1654
-
1655
- The branch is reset after each test that uses neon_branch_isolated.
1656
- """
1657
- env_manager = EnvironmentManager(_neon_config.env_var_name)
1658
- worker_id = _neon_xdist_coordinator.worker_id
1659
-
1660
- # Each worker creates its own isolated branch - no coordination needed
1661
- # because each worker has a unique ID
1662
- branch = _neon_branch_manager.create_branch(
1663
- name_suffix=f"-isolated-{worker_id}",
1664
- parent_branch_id=_neon_migration_branch.branch_id,
1665
- expiry_seconds=_neon_config.branch_expiry,
1666
- )
1667
-
1668
- # Store branch manager on config for reset operations
1669
- request.config._neon_isolated_branch_manager = _neon_branch_manager # type: ignore[attr-defined]
1670
-
1671
- # Set DATABASE_URL
1672
- env_manager.set(branch.connection_string)
1673
-
1674
- try:
1675
- yield branch
1676
- finally:
1677
- env_manager.restore()
1678
- _neon_branch_manager.delete_branch(branch.branch_id)
1679
-
1680
-
1681
- @pytest.fixture(scope="session")
1682
- def neon_branch_readonly(
1683
- _neon_config: NeonConfig,
1684
- _neon_readonly_endpoint: NeonBranch,
1685
955
  ) -> NeonBranch:
1686
956
  """
1687
- Provide a true read-only Neon database connection.
1688
-
1689
- This fixture uses a read_only endpoint on the migration branch, which
1690
- enforces read-only access at the database level. Any attempt to write
1691
- will result in a database error.
957
+ Provide a shared Neon database branch for all tests.
1692
958
 
1693
- This is the recommended fixture for tests that only read data (SELECT queries).
1694
- It's session-scoped and shared across all tests and workers since it's read-only.
959
+ This is a session-scoped branch that all tests share. Migrations run
960
+ once before tests start, then all tests see the same database state.
1695
961
 
1696
- Use this fixture when your tests only perform SELECT queries.
1697
- For tests that INSERT, UPDATE, or DELETE data, use ``neon_branch_dirty``
1698
- (for shared state) or ``neon_branch_isolated`` (for test isolation).
962
+ Since all tests share the branch:
963
+ - Data written by one test IS visible to subsequent tests
964
+ - Use transaction rollback or cleanup fixtures for test isolation
965
+ - Tests run in parallel (xdist) share the same branch
1699
966
 
1700
967
  The connection string is automatically set in the DATABASE_URL environment
1701
968
  variable (configurable via --neon-env-var).
@@ -1709,242 +976,44 @@ def neon_branch_readonly(
1709
976
 
1710
977
  Example::
1711
978
 
1712
- def test_query_users(neon_branch_readonly):
979
+ def test_query_users(neon_branch):
1713
980
  # DATABASE_URL is automatically set
1714
- conn_string = os.environ["DATABASE_URL"]
1715
-
1716
- # Read-only query
1717
- with psycopg.connect(conn_string) as conn:
981
+ import psycopg
982
+ with psycopg.connect(neon_branch.connection_string) as conn:
1718
983
  result = conn.execute("SELECT * FROM users").fetchall()
1719
- assert len(result) > 0
984
+ assert len(result) >= 0
1720
985
 
1721
- # This would fail with a database error:
1722
- # conn.execute("INSERT INTO users (name) VALUES ('test')")
1723
- """
1724
- # DATABASE_URL is already set by _neon_readonly_endpoint
1725
- return _neon_readonly_endpoint
986
+ For test isolation, use a transaction fixture::
1726
987
 
1727
-
1728
- @pytest.fixture(scope="session")
1729
- def neon_branch_dirty(
1730
- _neon_config: NeonConfig,
1731
- _neon_dirty_branch: NeonBranch,
1732
- ) -> NeonBranch:
1733
- """
1734
- Provide a session-scoped Neon database branch for read-write access.
1735
-
1736
- All tests share the same branch and writes persist across tests (no cleanup
1737
- between tests). This is faster than neon_branch_isolated because there's no
1738
- reset overhead.
1739
-
1740
- Use this fixture when:
1741
- - Most tests can share database state without interference
1742
- - You want maximum performance with minimal API calls
1743
- - You manually manage test data cleanup if needed
1744
- - You're using it alongside ``neon_branch_isolated`` for specific tests
1745
- that need guaranteed clean state
1746
-
1747
- The connection string is automatically set in the DATABASE_URL environment
1748
- variable (configurable via --neon-env-var).
1749
-
1750
- Warning:
1751
- Data written by one test WILL be visible to subsequent tests AND to
1752
- other xdist workers. This is truly shared - use ``neon_branch_isolated``
1753
- for tests that require guaranteed clean state.
1754
-
1755
- pytest-xdist:
1756
- ALL workers share the same dirty branch. Concurrent writes from different
1757
- workers may conflict. This is "dirty" by design - for isolation, use
1758
- ``neon_branch_isolated``.
1759
-
1760
- Requires either:
1761
- - NEON_API_KEY and NEON_PROJECT_ID environment variables, or
1762
- - --neon-api-key and --neon-project-id command line options
1763
-
1764
- Returns:
1765
- NeonBranch: Object with branch_id, project_id, connection_string, and host.
1766
-
1767
- Example::
1768
-
1769
- def test_insert_user(neon_branch_dirty):
1770
- # DATABASE_URL is automatically set
988
+ @pytest.fixture
989
+ def db_transaction(neon_branch):
1771
990
  import psycopg
1772
- with psycopg.connect(neon_branch_dirty.connection_string) as conn:
1773
- conn.execute("INSERT INTO users (name) VALUES ('test')")
1774
- conn.commit()
1775
- # Data persists - next test will see this user
1776
-
1777
- def test_count_users(neon_branch_dirty):
1778
- # This test sees data from previous tests
1779
- import psycopg
1780
- with psycopg.connect(neon_branch_dirty.connection_string) as conn:
1781
- result = conn.execute("SELECT COUNT(*) FROM users").fetchone()
1782
- # Count includes users from previous tests
1783
- """
1784
- # DATABASE_URL is already set by _neon_dirty_branch
1785
- return _neon_dirty_branch
1786
-
1787
-
1788
- @pytest.fixture(scope="function")
1789
- def neon_branch_isolated(
1790
- request: pytest.FixtureRequest,
1791
- _neon_config: NeonConfig,
1792
- _neon_isolated_branch: NeonBranch,
1793
- ) -> Generator[NeonBranch, None, None]:
991
+ conn = psycopg.connect(neon_branch.connection_string)
992
+ conn.execute("BEGIN")
993
+ yield conn
994
+ conn.execute("ROLLBACK")
995
+ conn.close()
996
+
997
+ def test_with_isolation(db_transaction):
998
+ db_transaction.execute("INSERT INTO users (name) VALUES ('test')")
999
+ # Rolled back after test - next test won't see this
1794
1000
  """
1795
- Provide an isolated Neon database branch with reset after each test.
1796
-
1797
- This is the recommended fixture for tests that modify database state and
1798
- need isolation. Each xdist worker has its own branch, and the branch is
1799
- reset to the migration state after each test.
1800
-
1801
- Use this fixture when:
1802
- - Tests modify database state (INSERT, UPDATE, DELETE)
1803
- - You need test isolation (each test starts with clean state)
1804
- - You're using it alongside ``neon_branch_dirty`` for specific tests
1805
-
1806
- The connection string is automatically set in the DATABASE_URL environment
1807
- variable (configurable via --neon-env-var).
1808
-
1809
- SQLAlchemy Users:
1810
- If you create your own engine (not using the neon_engine fixture),
1811
- you MUST use pool_pre_ping=True::
1001
+ branch, is_creator = _neon_test_branch
1812
1002
 
1813
- engine = create_engine(DATABASE_URL, pool_pre_ping=True)
1814
-
1815
- Branch resets terminate server-side connections. Without pool_pre_ping,
1816
- SQLAlchemy may reuse dead pooled connections, causing SSL errors.
1817
-
1818
- pytest-xdist:
1819
- Each worker has its own isolated branch. Resets only affect that worker's
1820
- branch, so workers don't interfere with each other.
1821
-
1822
- Requires either:
1823
- - NEON_API_KEY and NEON_PROJECT_ID environment variables, or
1824
- - --neon-api-key and --neon-project-id command line options
1825
-
1826
- Yields:
1827
- NeonBranch: Object with branch_id, project_id, connection_string, and host.
1828
-
1829
- Example::
1830
-
1831
- def test_insert_user(neon_branch_isolated):
1832
- # DATABASE_URL is automatically set
1833
- conn_string = os.environ["DATABASE_URL"]
1834
-
1835
- # Insert data - branch will reset after this test
1836
- with psycopg.connect(conn_string) as conn:
1837
- conn.execute("INSERT INTO users (name) VALUES ('test')")
1838
- conn.commit()
1839
- # Next test starts with clean state
1840
- """
1841
- # DATABASE_URL is already set by _neon_isolated_branch
1842
- yield _neon_isolated_branch
1843
-
1844
- # Reset branch to migration state after each test
1845
- branch_manager = getattr(request.config, "_neon_isolated_branch_manager", None)
1846
- if branch_manager is not None:
1847
- try:
1848
- branch_manager.reset_branch(_neon_isolated_branch)
1849
- except Exception as e:
1850
- pytest.fail(
1851
- f"\n\nFailed to reset branch {_neon_isolated_branch.branch_id} "
1852
- f"after test. Subsequent tests may see dirty state.\n\n"
1853
- f"Error: {e}\n\n"
1854
- f"To keep the branch for debugging, use --neon-keep-branches"
1855
- )
1856
-
1857
-
1858
- @pytest.fixture(scope="function")
1859
- def neon_branch_readwrite(
1860
- neon_branch_isolated: NeonBranch,
1861
- ) -> Generator[NeonBranch, None, None]:
1862
- """
1863
- Deprecated: Use ``neon_branch_isolated`` instead.
1864
-
1865
- This fixture is now an alias for ``neon_branch_isolated``.
1866
-
1867
- .. deprecated:: 2.3.0
1868
- Use ``neon_branch_isolated`` for tests that modify data with reset,
1869
- ``neon_branch_dirty`` for shared state, or ``neon_branch_readonly``
1870
- for read-only access.
1871
- """
1872
- warnings.warn(
1873
- "neon_branch_readwrite is deprecated. Use neon_branch_isolated (for tests "
1874
- "that modify data with isolation) or neon_branch_dirty (for shared state).",
1875
- DeprecationWarning,
1876
- stacklevel=2,
1877
- )
1878
- yield neon_branch_isolated
1879
-
1880
-
1881
- @pytest.fixture(scope="function")
1882
- def neon_branch(
1883
- neon_branch_isolated: NeonBranch,
1884
- ) -> Generator[NeonBranch, None, None]:
1885
- """
1886
- Deprecated: Use ``neon_branch_isolated``, ``neon_branch_dirty``, or
1887
- ``neon_branch_readonly`` instead.
1888
-
1889
- This fixture is now an alias for ``neon_branch_isolated``.
1890
-
1891
- .. deprecated:: 1.1.0
1892
- Use ``neon_branch_isolated`` for tests that modify data with reset,
1893
- ``neon_branch_dirty`` for shared state, or ``neon_branch_readonly``
1894
- for read-only access.
1895
- """
1896
- warnings.warn(
1897
- "neon_branch is deprecated. Use neon_branch_isolated (for tests that "
1898
- "modify data), neon_branch_dirty (for shared state), or "
1899
- "neon_branch_readonly (for read-only tests).",
1900
- DeprecationWarning,
1901
- stacklevel=2,
1902
- )
1903
- yield neon_branch_isolated
1904
-
1905
-
1906
- @pytest.fixture(scope="module")
1907
- def neon_branch_shared(
1908
- request: pytest.FixtureRequest,
1909
- _neon_migration_branch: NeonBranch,
1910
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1911
- ) -> Generator[NeonBranch, None, None]:
1912
- """
1913
- Provide a shared Neon database branch for all tests in a module.
1914
-
1915
- This fixture creates one branch per test module and shares it across all
1916
- tests without resetting. This is the fastest option but tests can see
1917
- each other's data modifications.
1918
-
1919
- If you override the `neon_apply_migrations` fixture, migrations will run
1920
- once before the first test, and this branch will include the migrated schema.
1921
-
1922
- Use this when:
1923
- - Tests are read-only or don't interfere with each other
1924
- - You manually clean up test data within each test
1925
- - Maximum speed is more important than isolation
1926
-
1927
- Warning: Tests in the same module will share database state. Data created
1928
- by one test will be visible to subsequent tests. Use `neon_branch` instead
1929
- if you need isolation between tests.
1930
-
1931
- Yields:
1932
- NeonBranch: Object with branch_id, project_id, connection_string, and host.
1003
+ if is_creator:
1004
+ # Creator runs migrations (via dependency), then signals completion
1005
+ _neon_xdist_coordinator.send_signal("migrations_done")
1006
+ else:
1007
+ # Non-creators wait for migrations to complete
1008
+ _neon_xdist_coordinator.wait_for_signal(
1009
+ "migrations_done", timeout=_MIGRATION_WAIT_TIMEOUT
1010
+ )
1933
1011
 
1934
- Example:
1935
- def test_read_only_query(neon_branch_shared):
1936
- # Fast: no reset between tests, but be careful about data leakage
1937
- conn_string = neon_branch_shared.connection_string
1938
- """
1939
- yield from _create_neon_branch(
1940
- request,
1941
- parent_branch_id_override=_neon_migration_branch.branch_id,
1942
- branch_name_suffix="-shared",
1943
- )
1012
+ return branch
1944
1013
 
1945
1014
 
1946
1015
  @pytest.fixture
1947
- def neon_connection(neon_branch_isolated: NeonBranch):
1016
+ def neon_connection(neon_branch: NeonBranch):
1948
1017
  """
1949
1018
  Provide a psycopg2 connection to the test branch.
1950
1019
 
@@ -1952,7 +1021,6 @@ def neon_connection(neon_branch_isolated: NeonBranch):
1952
1021
  pip install pytest-neon[psycopg2]
1953
1022
 
1954
1023
  The connection is rolled back and closed after each test.
1955
- Uses neon_branch_isolated for test isolation.
1956
1024
 
1957
1025
  Yields:
1958
1026
  psycopg2 connection object
@@ -1974,22 +1042,22 @@ def neon_connection(neon_branch_isolated: NeonBranch):
1974
1042
  " The 'neon_connection' fixture requires psycopg2.\n\n"
1975
1043
  " To fix this, install the psycopg2 extra:\n\n"
1976
1044
  " pip install pytest-neon[psycopg2]\n\n"
1977
- " Or use the 'neon_branch_isolated' fixture with your own driver:\n\n"
1978
- " def test_example(neon_branch_isolated):\n"
1045
+ " Or use the 'neon_branch' fixture with your own driver:\n\n"
1046
+ " def test_example(neon_branch):\n"
1979
1047
  " import your_driver\n"
1980
1048
  " conn = your_driver.connect(\n"
1981
- " neon_branch_isolated.connection_string)\n\n"
1049
+ " neon_branch.connection_string)\n\n"
1982
1050
  "═══════════════════════════════════════════════════════════════════\n"
1983
1051
  )
1984
1052
 
1985
- conn = psycopg2.connect(neon_branch_isolated.connection_string)
1053
+ conn = psycopg2.connect(neon_branch.connection_string)
1986
1054
  yield conn
1987
1055
  conn.rollback()
1988
1056
  conn.close()
1989
1057
 
1990
1058
 
1991
1059
  @pytest.fixture
1992
- def neon_connection_psycopg(neon_branch_isolated: NeonBranch):
1060
+ def neon_connection_psycopg(neon_branch: NeonBranch):
1993
1061
  """
1994
1062
  Provide a psycopg (v3) connection to the test branch.
1995
1063
 
@@ -1997,7 +1065,6 @@ def neon_connection_psycopg(neon_branch_isolated: NeonBranch):
1997
1065
  pip install pytest-neon[psycopg]
1998
1066
 
1999
1067
  The connection is rolled back and closed after each test.
2000
- Uses neon_branch_isolated for test isolation.
2001
1068
 
2002
1069
  Yields:
2003
1070
  psycopg connection object
@@ -2019,40 +1086,29 @@ def neon_connection_psycopg(neon_branch_isolated: NeonBranch):
2019
1086
  " The 'neon_connection_psycopg' fixture requires psycopg v3.\n\n"
2020
1087
  " To fix this, install the psycopg extra:\n\n"
2021
1088
  " pip install pytest-neon[psycopg]\n\n"
2022
- " Or use the 'neon_branch_isolated' fixture with your own driver:\n\n"
2023
- " def test_example(neon_branch_isolated):\n"
1089
+ " Or use the 'neon_branch' fixture with your own driver:\n\n"
1090
+ " def test_example(neon_branch):\n"
2024
1091
  " import your_driver\n"
2025
1092
  " conn = your_driver.connect(\n"
2026
- " neon_branch_isolated.connection_string)\n\n"
1093
+ " neon_branch.connection_string)\n\n"
2027
1094
  "═══════════════════════════════════════════════════════════════════\n"
2028
1095
  )
2029
1096
 
2030
- conn = psycopg.connect(neon_branch_isolated.connection_string)
1097
+ conn = psycopg.connect(neon_branch.connection_string)
2031
1098
  yield conn
2032
1099
  conn.rollback()
2033
1100
  conn.close()
2034
1101
 
2035
1102
 
2036
1103
  @pytest.fixture
2037
- def neon_engine(neon_branch_isolated: NeonBranch):
1104
+ def neon_engine(neon_branch: NeonBranch):
2038
1105
  """
2039
1106
  Provide a SQLAlchemy engine connected to the test branch.
2040
1107
 
2041
1108
  Requires the sqlalchemy optional dependency:
2042
1109
  pip install pytest-neon[sqlalchemy]
2043
1110
 
2044
- The engine is disposed after each test, which handles stale connections
2045
- after branch resets automatically. Uses neon_branch_isolated for test isolation.
2046
-
2047
- Note:
2048
- If you create your own module-level engine instead of using this
2049
- fixture, you MUST use pool_pre_ping=True::
2050
-
2051
- engine = create_engine(DATABASE_URL, pool_pre_ping=True)
2052
-
2053
- This is required because branch resets terminate server-side
2054
- connections, and without pool_pre_ping SQLAlchemy may reuse dead
2055
- pooled connections.
1111
+ The engine is disposed after each test.
2056
1112
 
2057
1113
  Yields:
2058
1114
  SQLAlchemy Engine object
@@ -2074,14 +1130,14 @@ def neon_engine(neon_branch_isolated: NeonBranch):
2074
1130
  " The 'neon_engine' fixture requires SQLAlchemy.\n\n"
2075
1131
  " To fix this, install the sqlalchemy extra:\n\n"
2076
1132
  " pip install pytest-neon[sqlalchemy]\n\n"
2077
- " Or use the 'neon_branch_isolated' fixture with your own driver:\n\n"
2078
- " def test_example(neon_branch_isolated):\n"
1133
+ " Or use the 'neon_branch' fixture with your own driver:\n\n"
1134
+ " def test_example(neon_branch):\n"
2079
1135
  " from sqlalchemy import create_engine\n"
2080
1136
  " engine = create_engine(\n"
2081
- " neon_branch_isolated.connection_string)\n\n"
1137
+ " neon_branch.connection_string)\n\n"
2082
1138
  "═══════════════════════════════════════════════════════════════════\n"
2083
1139
  )
2084
1140
 
2085
- engine = create_engine(neon_branch_isolated.connection_string)
1141
+ engine = create_engine(neon_branch.connection_string)
2086
1142
  yield engine
2087
1143
  engine.dispose()