pytest-neon 2.3.2__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,17 +290,6 @@ 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
293
  def _reveal_role_password(
310
294
  api_key: str, project_id: str, branch_id: str, role_name: str
311
295
  ) -> str:
@@ -337,38 +321,6 @@ def _reveal_role_password(
337
321
  return data["password"]
338
322
 
339
323
 
340
- def _get_schema_fingerprint(connection_string: str) -> tuple[tuple[Any, ...], ...]:
341
- """
342
- Get a fingerprint of the database schema for change detection.
343
-
344
- Queries information_schema for all tables, columns, and their properties
345
- in the public schema. Returns a hashable tuple that can be compared
346
- before/after migrations to detect if the schema actually changed.
347
-
348
- This is used to avoid creating unnecessary migration branches when
349
- no actual schema changes occurred.
350
- """
351
- try:
352
- import psycopg
353
- except ImportError:
354
- try:
355
- import psycopg2 as psycopg # type: ignore[import-not-found]
356
- except ImportError:
357
- # No driver available - can't fingerprint, assume migrations changed things
358
- return ()
359
-
360
- with psycopg.connect(connection_string) as conn, conn.cursor() as cur:
361
- cur.execute("""
362
- SELECT table_name, column_name, data_type, is_nullable,
363
- column_default, ordinal_position
364
- FROM information_schema.columns
365
- WHERE table_schema = 'public'
366
- ORDER BY table_name, ordinal_position
367
- """)
368
- rows = cur.fetchall()
369
- return tuple(tuple(row) for row in rows)
370
-
371
-
372
324
  @dataclass
373
325
  class NeonBranch:
374
326
  """Information about a Neon test branch."""
@@ -479,7 +431,7 @@ class NeonBranchManager:
479
431
  Create a new Neon branch with a read_write endpoint.
480
432
 
481
433
  Args:
482
- name_suffix: Suffix to add to branch name (e.g., "-migration", "-dirty")
434
+ name_suffix: Suffix to add to branch name (e.g., "-test")
483
435
  parent_branch_id: Parent branch ID (defaults to config's parent)
484
436
  expiry_seconds: Branch expiry in seconds (0 or None for no expiry)
485
437
 
@@ -552,53 +504,6 @@ class NeonBranchManager:
552
504
  endpoint_id=endpoint_id,
553
505
  )
554
506
 
555
- def create_readonly_endpoint(self, branch: NeonBranch) -> NeonBranch:
556
- """
557
- Create a read_only endpoint on an existing branch.
558
-
559
- This creates a true read-only endpoint that enforces no writes at the
560
- database level.
561
-
562
- Args:
563
- branch: The branch to create the endpoint on
564
-
565
- Returns:
566
- NeonBranch with the read_only endpoint's connection details
567
- """
568
- result = _retry_on_rate_limit(
569
- lambda: self._neon.endpoint_create(
570
- project_id=self.config.project_id,
571
- endpoint={
572
- "branch_id": branch.branch_id,
573
- "type": "read_only",
574
- },
575
- ),
576
- operation_name="endpoint_create_readonly",
577
- )
578
-
579
- endpoint_id = result.endpoint.id
580
- host = self._wait_for_endpoint(endpoint_id)
581
-
582
- # Reuse the password from the parent branch's connection string.
583
- # DO NOT call role_password_reset here - it would invalidate the
584
- # password used by the parent branch's read_write endpoint, breaking
585
- # any existing connections (especially in xdist where other workers
586
- # may be using the cached connection string).
587
- password = _extract_password_from_connection_string(branch.connection_string)
588
- connection_string = (
589
- f"postgresql://{self.config.role_name}:{password}@{host}/"
590
- f"{self.config.database_name}?sslmode=require"
591
- )
592
-
593
- return NeonBranch(
594
- branch_id=branch.branch_id,
595
- project_id=self.config.project_id,
596
- connection_string=connection_string,
597
- host=host,
598
- parent_id=branch.parent_id,
599
- endpoint_id=endpoint_id,
600
- )
601
-
602
507
  def delete_branch(self, branch_id: str) -> None:
603
508
  """Delete a branch (silently ignores errors)."""
604
509
  if self.config.keep_branches:
@@ -614,28 +519,6 @@ class NeonBranchManager:
614
519
  msg = f"Failed to delete Neon branch {branch_id}: {e}"
615
520
  warnings.warn(msg, stacklevel=2)
616
521
 
617
- def delete_endpoint(self, endpoint_id: str) -> None:
618
- """Delete an endpoint (silently ignores errors)."""
619
- try:
620
- _retry_on_rate_limit(
621
- lambda: self._neon.endpoint_delete(
622
- project_id=self.config.project_id, endpoint_id=endpoint_id
623
- ),
624
- operation_name="endpoint_delete",
625
- )
626
- except Exception as e:
627
- warnings.warn(
628
- f"Failed to delete Neon endpoint {endpoint_id}: {e}", stacklevel=2
629
- )
630
-
631
- def reset_branch(self, branch: NeonBranch) -> None:
632
- """Reset a branch to its parent's state."""
633
- if not branch.parent_id:
634
- msg = f"Branch {branch.branch_id} has no parent - cannot reset"
635
- raise RuntimeError(msg)
636
-
637
- _reset_branch_to_parent(branch, self.config.api_key)
638
-
639
522
  def _wait_for_endpoint(self, endpoint_id: str, max_wait_seconds: float = 60) -> str:
640
523
  """Wait for endpoint to become active and return its host."""
641
524
  poll_interval = 0.5
@@ -688,7 +571,7 @@ class XdistCoordinator:
688
571
  Coordinates branch sharing across pytest-xdist workers.
689
572
 
690
573
  Uses file locks and JSON cache files to ensure only one worker creates
691
- shared resources (like the migration branch), while others reuse them.
574
+ shared resources (like the test branch), while others reuse them.
692
575
  """
693
576
 
694
577
  def __init__(self, tmp_path_factory: pytest.TempPathFactory):
@@ -802,8 +685,7 @@ def _get_default_branch_id(neon: NeonAPI, project_id: str) -> str | None:
802
685
  Get the default/primary branch ID for a project.
803
686
 
804
687
  This is used as a safety check to ensure we never accidentally
805
- perform destructive operations (like password reset) on the
806
- production branch.
688
+ perform destructive operations on the production branch.
807
689
 
808
690
  Returns:
809
691
  The branch ID of the default branch, or None if not found.
@@ -932,447 +814,6 @@ def _get_config_value(
932
814
  return default
933
815
 
934
816
 
935
- def _create_neon_branch(
936
- request: pytest.FixtureRequest,
937
- parent_branch_id_override: str | None = None,
938
- branch_expiry_override: int | None = None,
939
- branch_name_suffix: str = "",
940
- ) -> Generator[NeonBranch, None, None]:
941
- """
942
- Internal helper that creates and manages a Neon branch lifecycle.
943
-
944
- This is the core implementation used by branch fixtures.
945
-
946
- Args:
947
- request: Pytest fixture request
948
- parent_branch_id_override: If provided, use this as parent instead of config
949
- branch_expiry_override: If provided, use this expiry instead of config
950
- branch_name_suffix: Optional suffix for branch name (e.g., "-migrated", "-test")
951
- """
952
- config = request.config
953
-
954
- api_key = _get_config_value(config, "neon_api_key", "NEON_API_KEY", "neon_api_key")
955
- project_id = _get_config_value(
956
- config, "neon_project_id", "NEON_PROJECT_ID", "neon_project_id"
957
- )
958
- # Use override if provided, otherwise read from config
959
- parent_branch_id = parent_branch_id_override or _get_config_value(
960
- config, "neon_parent_branch", "NEON_PARENT_BRANCH_ID", "neon_parent_branch"
961
- )
962
- database_name = _get_config_value(
963
- config, "neon_database", "NEON_DATABASE", "neon_database", "neondb"
964
- )
965
- role_name = _get_config_value(
966
- config, "neon_role", "NEON_ROLE", "neon_role", "neondb_owner"
967
- )
968
-
969
- # For boolean/int options, check CLI first, then ini
970
- keep_branches = config.getoption("neon_keep_branches", default=None)
971
- if keep_branches is None:
972
- keep_branches = config.getini("neon_keep_branches")
973
-
974
- # Use override if provided, otherwise read from config
975
- if branch_expiry_override is not None:
976
- branch_expiry = branch_expiry_override
977
- else:
978
- branch_expiry = config.getoption("neon_branch_expiry", default=None)
979
- if branch_expiry is None:
980
- branch_expiry = int(config.getini("neon_branch_expiry"))
981
-
982
- env_var_name = _get_config_value(
983
- config, "neon_env_var", "", "neon_env_var", "DATABASE_URL"
984
- )
985
-
986
- if not api_key:
987
- pytest.skip(
988
- "Neon API key not configured (set NEON_API_KEY or use --neon-api-key)"
989
- )
990
- if not project_id:
991
- pytest.skip(
992
- "Neon project ID not configured "
993
- "(set NEON_PROJECT_ID or use --neon-project-id)"
994
- )
995
-
996
- neon = NeonAPI(api_key=api_key)
997
-
998
- # Cache the default branch ID for safety checks (only fetch once per session)
999
- if not hasattr(config, "_neon_default_branch_id"):
1000
- config._neon_default_branch_id = _get_default_branch_id(neon, project_id) # type: ignore[attr-defined]
1001
-
1002
- # Generate unique branch name
1003
- # Format: pytest-[git branch (first 15 chars)]-[random]-[suffix]
1004
- # This helps identify orphaned branches by showing which git branch created them
1005
- random_suffix = os.urandom(2).hex() # 2 bytes = 4 hex chars
1006
- git_branch = _get_git_branch_name()
1007
- if git_branch:
1008
- # Truncate git branch to 15 chars to keep branch names reasonable
1009
- git_prefix = git_branch[:15]
1010
- branch_name = f"pytest-{git_prefix}-{random_suffix}{branch_name_suffix}"
1011
- else:
1012
- branch_name = f"pytest-{random_suffix}{branch_name_suffix}"
1013
-
1014
- # Build branch creation payload
1015
- branch_config: dict[str, Any] = {"name": branch_name}
1016
- if parent_branch_id:
1017
- branch_config["parent_id"] = parent_branch_id
1018
-
1019
- # Set branch expiration (auto-delete) as a safety net for interrupted test runs
1020
- # This uses the branch expires_at field, not endpoint suspend_timeout
1021
- if branch_expiry and branch_expiry > 0:
1022
- expires_at = datetime.now(timezone.utc) + timedelta(seconds=branch_expiry)
1023
- branch_config["expires_at"] = expires_at.strftime("%Y-%m-%dT%H:%M:%SZ")
1024
-
1025
- # Create branch with compute endpoint
1026
- # Wrap in retry logic to handle rate limits
1027
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1028
- result = _retry_on_rate_limit(
1029
- lambda: neon.branch_create(
1030
- project_id=project_id,
1031
- branch=branch_config,
1032
- endpoints=[{"type": "read_write"}],
1033
- ),
1034
- operation_name="branch_create",
1035
- )
1036
-
1037
- branch = result.branch
1038
-
1039
- # Get endpoint_id from operations
1040
- # (branch_create returns operations, not endpoints directly)
1041
- endpoint_id = None
1042
- for op in result.operations:
1043
- if op.endpoint_id:
1044
- endpoint_id = op.endpoint_id
1045
- break
1046
-
1047
- if not endpoint_id:
1048
- raise RuntimeError(f"No endpoint created for branch {branch.id}")
1049
-
1050
- # Wait for endpoint to be ready (it starts in "init" state)
1051
- # Endpoints typically become active in 1-2 seconds, but we allow up to 60s
1052
- # to handle occasional Neon API slowness or high load scenarios
1053
- max_wait_seconds = 60
1054
- poll_interval = 0.5 # Poll every 500ms for responsive feedback
1055
- waited = 0.0
1056
-
1057
- while True:
1058
- # Wrap in retry logic to handle rate limits during polling
1059
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1060
- endpoint_response = _retry_on_rate_limit(
1061
- lambda: neon.endpoint(project_id=project_id, endpoint_id=endpoint_id),
1062
- operation_name="endpoint_status",
1063
- )
1064
- endpoint = endpoint_response.endpoint
1065
- state = endpoint.current_state
1066
-
1067
- if state == EndpointState.active:
1068
- break
1069
-
1070
- if waited >= max_wait_seconds:
1071
- raise RuntimeError(
1072
- f"Timeout waiting for endpoint {endpoint_id} to become active "
1073
- f"(current state: {state})"
1074
- )
1075
-
1076
- time.sleep(poll_interval)
1077
- waited += poll_interval
1078
-
1079
- host = endpoint.host
1080
-
1081
- # Get password using reveal (not reset) to avoid invalidating existing connections
1082
- # See: https://api-docs.neon.tech/reference/getprojectbranchrolepassword
1083
- password = _retry_on_rate_limit(
1084
- lambda: _reveal_role_password(
1085
- api_key=api_key,
1086
- project_id=project_id,
1087
- branch_id=branch.id,
1088
- role_name=role_name,
1089
- ),
1090
- operation_name="role_password_reveal",
1091
- )
1092
-
1093
- # Build connection string
1094
- connection_string = (
1095
- f"postgresql://{role_name}:{password}@{host}/{database_name}?sslmode=require"
1096
- )
1097
-
1098
- neon_branch_info = NeonBranch(
1099
- branch_id=branch.id,
1100
- project_id=project_id,
1101
- connection_string=connection_string,
1102
- host=host,
1103
- parent_id=branch.parent_id,
1104
- endpoint_id=endpoint_id,
1105
- )
1106
-
1107
- # Set DATABASE_URL (or configured env var) for the duration of the fixture scope
1108
- original_env_value = os.environ.get(env_var_name)
1109
- os.environ[env_var_name] = connection_string
1110
-
1111
- try:
1112
- yield neon_branch_info
1113
- finally:
1114
- # Restore original env var
1115
- if original_env_value is None:
1116
- os.environ.pop(env_var_name, None)
1117
- else:
1118
- os.environ[env_var_name] = original_env_value
1119
-
1120
- # Cleanup: delete branch unless --neon-keep-branches was specified
1121
- if not keep_branches:
1122
- try:
1123
- # Wrap in retry logic to handle rate limits
1124
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1125
- _retry_on_rate_limit(
1126
- lambda: neon.branch_delete(
1127
- project_id=project_id, branch_id=branch.id
1128
- ),
1129
- operation_name="branch_delete",
1130
- )
1131
- except Exception as e:
1132
- # Log but don't fail tests due to cleanup issues
1133
- warnings.warn(
1134
- f"Failed to delete Neon branch {branch.id}: {e}",
1135
- stacklevel=2,
1136
- )
1137
-
1138
-
1139
- def _create_readonly_endpoint(
1140
- branch: NeonBranch,
1141
- api_key: str,
1142
- database_name: str,
1143
- role_name: str,
1144
- ) -> NeonBranch:
1145
- """
1146
- Create a read_only endpoint on an existing branch.
1147
-
1148
- Returns a new NeonBranch object with the read_only endpoint's connection string.
1149
- The read_only endpoint enforces that no writes can be made through this connection.
1150
-
1151
- Args:
1152
- branch: The branch to create a read_only endpoint on
1153
- api_key: Neon API key
1154
- database_name: Database name for connection string
1155
- role_name: Role name for connection string
1156
-
1157
- Returns:
1158
- NeonBranch with read_only endpoint connection details
1159
- """
1160
- neon = NeonAPI(api_key=api_key)
1161
-
1162
- # Create read_only endpoint on the branch
1163
- # See: https://api-docs.neon.tech/reference/createprojectendpoint
1164
- result = _retry_on_rate_limit(
1165
- lambda: neon.endpoint_create(
1166
- project_id=branch.project_id,
1167
- endpoint={
1168
- "branch_id": branch.branch_id,
1169
- "type": "read_only",
1170
- },
1171
- ),
1172
- operation_name="endpoint_create_readonly",
1173
- )
1174
-
1175
- endpoint = result.endpoint
1176
- endpoint_id = endpoint.id
1177
-
1178
- # Wait for endpoint to be ready
1179
- max_wait_seconds = 60
1180
- poll_interval = 0.5
1181
- waited = 0.0
1182
-
1183
- while True:
1184
- endpoint_response = _retry_on_rate_limit(
1185
- lambda: neon.endpoint(
1186
- project_id=branch.project_id, endpoint_id=endpoint_id
1187
- ),
1188
- operation_name="endpoint_status_readonly",
1189
- )
1190
- endpoint = endpoint_response.endpoint
1191
- state = endpoint.current_state
1192
-
1193
- if state == EndpointState.active:
1194
- break
1195
-
1196
- if waited >= max_wait_seconds:
1197
- raise RuntimeError(
1198
- f"Timeout waiting for read_only endpoint {endpoint_id} "
1199
- f"to become active (current state: {state})"
1200
- )
1201
-
1202
- time.sleep(poll_interval)
1203
- waited += poll_interval
1204
-
1205
- host = endpoint.host
1206
-
1207
- # Reuse the password from the parent branch's connection string.
1208
- # DO NOT call role_password_reset here - it would invalidate the
1209
- # password used by the parent branch's read_write endpoint.
1210
- password = _extract_password_from_connection_string(branch.connection_string)
1211
-
1212
- # Build connection string for the read_only endpoint
1213
- connection_string = (
1214
- f"postgresql://{role_name}:{password}@{host}/{database_name}?sslmode=require"
1215
- )
1216
-
1217
- return NeonBranch(
1218
- branch_id=branch.branch_id,
1219
- project_id=branch.project_id,
1220
- connection_string=connection_string,
1221
- host=host,
1222
- parent_id=branch.parent_id,
1223
- endpoint_id=endpoint_id,
1224
- )
1225
-
1226
-
1227
- def _delete_endpoint(project_id: str, endpoint_id: str, api_key: str) -> None:
1228
- """Delete a Neon endpoint."""
1229
- neon = NeonAPI(api_key=api_key)
1230
- try:
1231
- _retry_on_rate_limit(
1232
- lambda: neon.endpoint_delete(
1233
- project_id=project_id, endpoint_id=endpoint_id
1234
- ),
1235
- operation_name="endpoint_delete",
1236
- )
1237
- except Exception as e:
1238
- warnings.warn(
1239
- f"Failed to delete Neon endpoint {endpoint_id}: {e}",
1240
- stacklevel=2,
1241
- )
1242
-
1243
-
1244
- def _reset_branch_to_parent(branch: NeonBranch, api_key: str) -> None:
1245
- """Reset a branch to its parent's state using the Neon API.
1246
-
1247
- Uses exponential backoff retry logic with jitter to handle rate limit (429)
1248
- errors. After initiating the restore, polls the operation status until it
1249
- completes.
1250
-
1251
- See: https://api-docs.neon.tech/reference/api-rate-limiting
1252
-
1253
- Args:
1254
- branch: The branch to reset
1255
- api_key: Neon API key
1256
- """
1257
- if not branch.parent_id:
1258
- raise RuntimeError(f"Branch {branch.branch_id} has no parent - cannot reset")
1259
-
1260
- base_url = "https://console.neon.tech/api/v2"
1261
- project_id = branch.project_id
1262
- branch_id = branch.branch_id
1263
- restore_url = f"{base_url}/projects/{project_id}/branches/{branch_id}/restore"
1264
- headers = {
1265
- "Authorization": f"Bearer {api_key}",
1266
- "Content-Type": "application/json",
1267
- }
1268
-
1269
- def do_restore() -> dict[str, Any]:
1270
- response = requests.post(
1271
- restore_url,
1272
- headers=headers,
1273
- json={"source_branch_id": branch.parent_id},
1274
- timeout=30,
1275
- )
1276
- response.raise_for_status()
1277
- return response.json()
1278
-
1279
- # Wrap in retry logic to handle rate limits
1280
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1281
- data = _retry_on_rate_limit(do_restore, operation_name="branch_restore")
1282
- operations = data.get("operations", [])
1283
-
1284
- # The restore API returns operations that run asynchronously.
1285
- # We must wait for operations to complete before the next test
1286
- # starts, otherwise connections may fail during the restore.
1287
- if operations:
1288
- _wait_for_operations(
1289
- project_id=branch.project_id,
1290
- operations=operations,
1291
- headers=headers,
1292
- base_url=base_url,
1293
- )
1294
-
1295
-
1296
- def _wait_for_operations(
1297
- project_id: str,
1298
- operations: list[dict[str, Any]],
1299
- headers: dict[str, str],
1300
- base_url: str,
1301
- max_wait_seconds: float = 60,
1302
- poll_interval: float = 0.5,
1303
- ) -> None:
1304
- """Wait for Neon operations to complete.
1305
-
1306
- Handles rate limit (429) errors with exponential backoff retry.
1307
- See: https://api-docs.neon.tech/reference/api-rate-limiting
1308
-
1309
- Args:
1310
- project_id: The Neon project ID
1311
- operations: List of operation dicts from the API response
1312
- headers: HTTP headers including auth
1313
- base_url: Base URL for Neon API
1314
- max_wait_seconds: Maximum time to wait (default: 60s)
1315
- poll_interval: Time between polls (default: 0.5s)
1316
- """
1317
- # Get operation IDs that aren't already finished
1318
- pending_op_ids = [
1319
- op["id"] for op in operations if op.get("status") not in ("finished", "skipped")
1320
- ]
1321
-
1322
- if not pending_op_ids:
1323
- return # All operations already complete
1324
-
1325
- waited = 0.0
1326
- first_poll = True
1327
- while pending_op_ids and waited < max_wait_seconds:
1328
- # Poll immediately first time (operation usually completes instantly),
1329
- # then wait between subsequent polls
1330
- if first_poll:
1331
- time.sleep(0.1) # Tiny delay to let operation start
1332
- waited += 0.1
1333
- first_poll = False
1334
- else:
1335
- time.sleep(poll_interval)
1336
- waited += poll_interval
1337
-
1338
- # Check status of each pending operation
1339
- still_pending = []
1340
- for op_id in pending_op_ids:
1341
- op_url = f"{base_url}/projects/{project_id}/operations/{op_id}"
1342
-
1343
- def get_operation_status(url: str = op_url) -> dict[str, Any]:
1344
- """Fetch operation status. Default arg captures url by value."""
1345
- response = requests.get(url, headers=headers, timeout=10)
1346
- response.raise_for_status()
1347
- return response.json()
1348
-
1349
- try:
1350
- # Wrap in retry logic to handle rate limits
1351
- # See: https://api-docs.neon.tech/reference/api-rate-limiting
1352
- result = _retry_on_rate_limit(
1353
- get_operation_status,
1354
- operation_name=f"operation_status({op_id})",
1355
- )
1356
- op_data = result.get("operation", {})
1357
- status = op_data.get("status")
1358
-
1359
- if status == "failed":
1360
- err = op_data.get("error", "unknown error")
1361
- raise RuntimeError(f"Operation {op_id} failed: {err}")
1362
- if status not in ("finished", "skipped", "cancelled"):
1363
- still_pending.append(op_id)
1364
- except requests.RequestException:
1365
- # On network error (non-429), assume still pending and retry
1366
- still_pending.append(op_id)
1367
-
1368
- pending_op_ids = still_pending
1369
-
1370
- if pending_op_ids:
1371
- raise RuntimeError(
1372
- f"Timeout waiting for operations to complete: {pending_op_ids}"
1373
- )
1374
-
1375
-
1376
817
  def _branch_to_dict(branch: NeonBranch) -> dict[str, Any]:
1377
818
  """Convert NeonBranch to a JSON-serializable dict."""
1378
819
  return asdict(branch)
@@ -1418,301 +859,110 @@ def _neon_xdist_coordinator(
1418
859
 
1419
860
 
1420
861
  @pytest.fixture(scope="session")
1421
- def _neon_migration_branch(
1422
- request: pytest.FixtureRequest,
862
+ def _neon_test_branch(
1423
863
  _neon_config: NeonConfig,
1424
864
  _neon_branch_manager: NeonBranchManager,
1425
865
  _neon_xdist_coordinator: XdistCoordinator,
1426
- ) -> Generator[NeonBranch, None, None]:
866
+ ) -> Generator[tuple[NeonBranch, bool], None, None]:
1427
867
  """
1428
- Session-scoped branch where migrations are applied.
868
+ Internal: Create test branch, coordinated across workers.
1429
869
 
1430
- This branch is ALWAYS created from the configured parent and serves as
1431
- the parent for all test branches (dirty, isolated, readonly endpoint).
1432
- 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.
1433
872
 
1434
- pytest-xdist Support:
1435
- When running with pytest-xdist, the first worker to acquire the lock
1436
- creates the migration branch. Other workers wait for migrations to
1437
- complete, then reuse the same branch. This avoids redundant API calls
1438
- and ensures migrations only run once. Only the creator cleans up the
1439
- branch at session end.
1440
-
1441
- Note: The migration branch cannot have an expiry because Neon doesn't
1442
- allow creating child branches from branches with expiration dates.
1443
- 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).
1444
876
  """
1445
877
  env_manager = EnvironmentManager(_neon_config.env_var_name)
1446
- branch: NeonBranch
1447
- is_creator: bool
1448
878
 
1449
- def create_migration_branch() -> dict[str, Any]:
879
+ def create_branch() -> dict[str, Any]:
1450
880
  b = _neon_branch_manager.create_branch(
1451
- name_suffix="-migration",
1452
- expiry_seconds=0, # No expiry - child branches need this
881
+ name_suffix="-test",
882
+ expiry_seconds=_neon_config.branch_expiry,
1453
883
  )
1454
884
  return {"branch": _branch_to_dict(b)}
1455
885
 
1456
- # Coordinate branch creation across xdist workers
1457
886
  data, is_creator = _neon_xdist_coordinator.coordinate_resource(
1458
- "migration_branch", create_migration_branch
887
+ "test_branch", create_branch
1459
888
  )
1460
889
  branch = _dict_to_branch(data["branch"])
1461
-
1462
- # Store creator status for other fixtures
1463
- request.config._neon_is_migration_creator = is_creator # type: ignore[attr-defined]
1464
-
1465
- # Set DATABASE_URL
1466
890
  env_manager.set(branch.connection_string)
1467
891
 
1468
- # Non-creators wait for migrations to complete
1469
- if not is_creator:
1470
- _neon_xdist_coordinator.wait_for_signal(
1471
- "migrations_done", timeout=_MIGRATION_WAIT_TIMEOUT
1472
- )
1473
-
1474
892
  try:
1475
- yield branch
893
+ yield branch, is_creator
1476
894
  finally:
1477
895
  env_manager.restore()
1478
- # Only creator cleans up
1479
896
  if is_creator:
1480
897
  _neon_branch_manager.delete_branch(branch.branch_id)
1481
898
 
1482
899
 
1483
900
  @pytest.fixture(scope="session")
1484
- def neon_apply_migrations(_neon_migration_branch: NeonBranch) -> Any:
901
+ def neon_apply_migrations(_neon_test_branch: tuple[NeonBranch, bool]) -> Any:
1485
902
  """
1486
903
  Override this fixture to run migrations on the test database.
1487
904
 
1488
- The migration branch is already created and DATABASE_URL is set.
905
+ The test branch is already created and DATABASE_URL is set.
1489
906
  Migrations run once per test session, before any tests execute.
1490
907
 
1491
908
  pytest-xdist Support:
1492
909
  When running with pytest-xdist, migrations only run on the first
1493
- worker (the one that created the migration branch). Other workers
910
+ worker (the one that created the test branch). Other workers
1494
911
  wait for migrations to complete before proceeding. This ensures
1495
912
  migrations run exactly once, even with parallel workers.
1496
913
 
1497
- Smart Migration Detection:
1498
- The plugin automatically detects whether migrations actually modified
1499
- the database schema. If no schema changes occurred (or this fixture
1500
- isn't overridden), the plugin skips creating a separate migration
1501
- branch, saving Neon costs and branch slots.
1502
-
1503
914
  Example in conftest.py:
1504
915
 
1505
916
  @pytest.fixture(scope="session")
1506
- def neon_apply_migrations(_neon_migration_branch):
917
+ def neon_apply_migrations(_neon_test_branch):
1507
918
  import subprocess
1508
919
  subprocess.run(["alembic", "upgrade", "head"], check=True)
1509
920
 
1510
921
  Or with Django:
1511
922
 
1512
923
  @pytest.fixture(scope="session")
1513
- def neon_apply_migrations(_neon_migration_branch):
924
+ def neon_apply_migrations(_neon_test_branch):
1514
925
  from django.core.management import call_command
1515
926
  call_command("migrate", "--noinput")
1516
927
 
1517
928
  Or with raw SQL:
1518
929
 
1519
930
  @pytest.fixture(scope="session")
1520
- def neon_apply_migrations(_neon_migration_branch):
931
+ def neon_apply_migrations(_neon_test_branch):
1521
932
  import psycopg
1522
- 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:
1523
935
  with open("schema.sql") as f:
1524
936
  conn.execute(f.read())
1525
937
  conn.commit()
1526
938
 
1527
939
  Args:
1528
- _neon_migration_branch: The migration branch with connection details.
1529
- 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,
1530
942
  or rely on DATABASE_URL which is already set.
1531
943
 
1532
944
  Returns:
1533
- Any value (ignored). The default returns a sentinel to indicate
1534
- the fixture was not overridden.
945
+ Any value (ignored). The default returns None.
1535
946
  """
1536
- return _MIGRATIONS_NOT_DEFINED
947
+ return None
1537
948
 
1538
949
 
1539
950
  @pytest.fixture(scope="session")
1540
- def _neon_migrations_synchronized(
1541
- request: pytest.FixtureRequest,
1542
- _neon_migration_branch: NeonBranch,
951
+ def neon_branch(
952
+ _neon_test_branch: tuple[NeonBranch, bool],
1543
953
  _neon_xdist_coordinator: XdistCoordinator,
1544
954
  neon_apply_migrations: Any,
1545
- ) -> Any:
1546
- """
1547
- Internal fixture that synchronizes migrations across xdist workers.
1548
-
1549
- This fixture ensures that:
1550
- 1. Only the creator worker runs migrations (non-creators wait in
1551
- _neon_migration_branch BEFORE neon_apply_migrations runs)
1552
- 2. Creator signals completion after migrations finish
1553
- 3. The return value from neon_apply_migrations is preserved for detection
1554
-
1555
- Without xdist, this is a simple passthrough.
1556
- """
1557
- is_creator = getattr(request.config, "_neon_is_migration_creator", True)
1558
-
1559
- if is_creator:
1560
- # Creator: migrations just ran via neon_apply_migrations dependency
1561
- # Signal completion to other workers
1562
- _neon_xdist_coordinator.send_signal("migrations_done")
1563
-
1564
- return neon_apply_migrations
1565
-
1566
-
1567
- @pytest.fixture(scope="session")
1568
- def _neon_dirty_branch(
1569
- _neon_config: NeonConfig,
1570
- _neon_branch_manager: NeonBranchManager,
1571
- _neon_xdist_coordinator: XdistCoordinator,
1572
- _neon_migration_branch: NeonBranch,
1573
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1574
- ) -> Generator[NeonBranch, None, None]:
1575
- """
1576
- Session-scoped dirty branch shared across ALL xdist workers.
1577
-
1578
- This branch is a child of the migration branch. All tests using
1579
- neon_branch_dirty share this single branch - writes persist and
1580
- are visible to all tests (even across workers).
1581
-
1582
- This is the "dirty" branch because:
1583
- - No reset between tests
1584
- - Shared across all workers (concurrent writes possible)
1585
- - Fast because no per-test overhead
1586
- """
1587
- env_manager = EnvironmentManager(_neon_config.env_var_name)
1588
- branch: NeonBranch
1589
- is_creator: bool
1590
-
1591
- def create_dirty_branch() -> dict[str, Any]:
1592
- b = _neon_branch_manager.create_branch(
1593
- name_suffix="-dirty",
1594
- parent_branch_id=_neon_migration_branch.branch_id,
1595
- expiry_seconds=_neon_config.branch_expiry,
1596
- )
1597
- return {"branch": _branch_to_dict(b)}
1598
-
1599
- # Coordinate dirty branch creation - shared across ALL workers
1600
- data, is_creator = _neon_xdist_coordinator.coordinate_resource(
1601
- "dirty_branch", create_dirty_branch
1602
- )
1603
- branch = _dict_to_branch(data["branch"])
1604
-
1605
- # Set DATABASE_URL
1606
- env_manager.set(branch.connection_string)
1607
-
1608
- try:
1609
- yield branch
1610
- finally:
1611
- env_manager.restore()
1612
- # Only creator cleans up
1613
- if is_creator:
1614
- _neon_branch_manager.delete_branch(branch.branch_id)
1615
-
1616
-
1617
- @pytest.fixture(scope="session")
1618
- def _neon_readonly_endpoint(
1619
- _neon_config: NeonConfig,
1620
- _neon_branch_manager: NeonBranchManager,
1621
- _neon_xdist_coordinator: XdistCoordinator,
1622
- _neon_migration_branch: NeonBranch,
1623
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1624
- ) -> Generator[NeonBranch, None, None]:
1625
- """
1626
- Session-scoped read_only endpoint on the migration branch.
1627
-
1628
- This is a true read-only endpoint - writes are blocked at the database
1629
- level. All workers share this endpoint since it's read-only anyway.
1630
- """
1631
- env_manager = EnvironmentManager(_neon_config.env_var_name)
1632
- branch: NeonBranch
1633
- is_creator: bool
1634
-
1635
- def create_readonly_endpoint() -> dict[str, Any]:
1636
- b = _neon_branch_manager.create_readonly_endpoint(_neon_migration_branch)
1637
- return {"branch": _branch_to_dict(b)}
1638
-
1639
- # Coordinate endpoint creation - shared across ALL workers
1640
- data, is_creator = _neon_xdist_coordinator.coordinate_resource(
1641
- "readonly_endpoint", create_readonly_endpoint
1642
- )
1643
- branch = _dict_to_branch(data["branch"])
1644
-
1645
- # Set DATABASE_URL
1646
- env_manager.set(branch.connection_string)
1647
-
1648
- try:
1649
- yield branch
1650
- finally:
1651
- env_manager.restore()
1652
- # Only creator cleans up the endpoint
1653
- if is_creator and branch.endpoint_id:
1654
- _neon_branch_manager.delete_endpoint(branch.endpoint_id)
1655
-
1656
-
1657
- @pytest.fixture(scope="session")
1658
- def _neon_isolated_branch(
1659
- request: pytest.FixtureRequest,
1660
- _neon_config: NeonConfig,
1661
- _neon_branch_manager: NeonBranchManager,
1662
- _neon_xdist_coordinator: XdistCoordinator,
1663
- _neon_migration_branch: NeonBranch,
1664
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1665
- ) -> Generator[NeonBranch, None, None]:
1666
- """
1667
- Session-scoped isolated branch, one per xdist worker.
1668
-
1669
- Each worker gets its own branch. Unlike the dirty branch, this is
1670
- per-worker to allow reset operations without affecting other workers.
1671
-
1672
- The branch is reset after each test that uses neon_branch_isolated.
1673
- """
1674
- env_manager = EnvironmentManager(_neon_config.env_var_name)
1675
- worker_id = _neon_xdist_coordinator.worker_id
1676
-
1677
- # Each worker creates its own isolated branch - no coordination needed
1678
- # because each worker has a unique ID
1679
- branch = _neon_branch_manager.create_branch(
1680
- name_suffix=f"-isolated-{worker_id}",
1681
- parent_branch_id=_neon_migration_branch.branch_id,
1682
- expiry_seconds=_neon_config.branch_expiry,
1683
- )
1684
-
1685
- # Store branch manager on config for reset operations
1686
- request.config._neon_isolated_branch_manager = _neon_branch_manager # type: ignore[attr-defined]
1687
-
1688
- # Set DATABASE_URL
1689
- env_manager.set(branch.connection_string)
1690
-
1691
- try:
1692
- yield branch
1693
- finally:
1694
- env_manager.restore()
1695
- _neon_branch_manager.delete_branch(branch.branch_id)
1696
-
1697
-
1698
- @pytest.fixture(scope="session")
1699
- def neon_branch_readonly(
1700
- _neon_config: NeonConfig,
1701
- _neon_readonly_endpoint: NeonBranch,
1702
955
  ) -> NeonBranch:
1703
956
  """
1704
- Provide a true read-only Neon database connection.
1705
-
1706
- This fixture uses a read_only endpoint on the migration branch, which
1707
- enforces read-only access at the database level. Any attempt to write
1708
- will result in a database error.
957
+ Provide a shared Neon database branch for all tests.
1709
958
 
1710
- This is the recommended fixture for tests that only read data (SELECT queries).
1711
- 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.
1712
961
 
1713
- Use this fixture when your tests only perform SELECT queries.
1714
- For tests that INSERT, UPDATE, or DELETE data, use ``neon_branch_dirty``
1715
- (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
1716
966
 
1717
967
  The connection string is automatically set in the DATABASE_URL environment
1718
968
  variable (configurable via --neon-env-var).
@@ -1726,242 +976,44 @@ def neon_branch_readonly(
1726
976
 
1727
977
  Example::
1728
978
 
1729
- def test_query_users(neon_branch_readonly):
979
+ def test_query_users(neon_branch):
1730
980
  # DATABASE_URL is automatically set
1731
- conn_string = os.environ["DATABASE_URL"]
1732
-
1733
- # Read-only query
1734
- with psycopg.connect(conn_string) as conn:
981
+ import psycopg
982
+ with psycopg.connect(neon_branch.connection_string) as conn:
1735
983
  result = conn.execute("SELECT * FROM users").fetchall()
1736
- assert len(result) > 0
984
+ assert len(result) >= 0
1737
985
 
1738
- # This would fail with a database error:
1739
- # conn.execute("INSERT INTO users (name) VALUES ('test')")
1740
- """
1741
- # DATABASE_URL is already set by _neon_readonly_endpoint
1742
- return _neon_readonly_endpoint
986
+ For test isolation, use a transaction fixture::
1743
987
 
1744
-
1745
- @pytest.fixture(scope="session")
1746
- def neon_branch_dirty(
1747
- _neon_config: NeonConfig,
1748
- _neon_dirty_branch: NeonBranch,
1749
- ) -> NeonBranch:
1750
- """
1751
- Provide a session-scoped Neon database branch for read-write access.
1752
-
1753
- All tests share the same branch and writes persist across tests (no cleanup
1754
- between tests). This is faster than neon_branch_isolated because there's no
1755
- reset overhead.
1756
-
1757
- Use this fixture when:
1758
- - Most tests can share database state without interference
1759
- - You want maximum performance with minimal API calls
1760
- - You manually manage test data cleanup if needed
1761
- - You're using it alongside ``neon_branch_isolated`` for specific tests
1762
- that need guaranteed clean state
1763
-
1764
- The connection string is automatically set in the DATABASE_URL environment
1765
- variable (configurable via --neon-env-var).
1766
-
1767
- Warning:
1768
- Data written by one test WILL be visible to subsequent tests AND to
1769
- other xdist workers. This is truly shared - use ``neon_branch_isolated``
1770
- for tests that require guaranteed clean state.
1771
-
1772
- pytest-xdist:
1773
- ALL workers share the same dirty branch. Concurrent writes from different
1774
- workers may conflict. This is "dirty" by design - for isolation, use
1775
- ``neon_branch_isolated``.
1776
-
1777
- Requires either:
1778
- - NEON_API_KEY and NEON_PROJECT_ID environment variables, or
1779
- - --neon-api-key and --neon-project-id command line options
1780
-
1781
- Returns:
1782
- NeonBranch: Object with branch_id, project_id, connection_string, and host.
1783
-
1784
- Example::
1785
-
1786
- def test_insert_user(neon_branch_dirty):
1787
- # DATABASE_URL is automatically set
988
+ @pytest.fixture
989
+ def db_transaction(neon_branch):
1788
990
  import psycopg
1789
- with psycopg.connect(neon_branch_dirty.connection_string) as conn:
1790
- conn.execute("INSERT INTO users (name) VALUES ('test')")
1791
- conn.commit()
1792
- # Data persists - next test will see this user
1793
-
1794
- def test_count_users(neon_branch_dirty):
1795
- # This test sees data from previous tests
1796
- import psycopg
1797
- with psycopg.connect(neon_branch_dirty.connection_string) as conn:
1798
- result = conn.execute("SELECT COUNT(*) FROM users").fetchone()
1799
- # Count includes users from previous tests
1800
- """
1801
- # DATABASE_URL is already set by _neon_dirty_branch
1802
- return _neon_dirty_branch
1803
-
1804
-
1805
- @pytest.fixture(scope="function")
1806
- def neon_branch_isolated(
1807
- request: pytest.FixtureRequest,
1808
- _neon_config: NeonConfig,
1809
- _neon_isolated_branch: NeonBranch,
1810
- ) -> 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
1811
1000
  """
1812
- Provide an isolated Neon database branch with reset after each test.
1813
-
1814
- This is the recommended fixture for tests that modify database state and
1815
- need isolation. Each xdist worker has its own branch, and the branch is
1816
- reset to the migration state after each test.
1817
-
1818
- Use this fixture when:
1819
- - Tests modify database state (INSERT, UPDATE, DELETE)
1820
- - You need test isolation (each test starts with clean state)
1821
- - You're using it alongside ``neon_branch_dirty`` for specific tests
1822
-
1823
- The connection string is automatically set in the DATABASE_URL environment
1824
- variable (configurable via --neon-env-var).
1825
-
1826
- SQLAlchemy Users:
1827
- If you create your own engine (not using the neon_engine fixture),
1828
- you MUST use pool_pre_ping=True::
1001
+ branch, is_creator = _neon_test_branch
1829
1002
 
1830
- engine = create_engine(DATABASE_URL, pool_pre_ping=True)
1831
-
1832
- Branch resets terminate server-side connections. Without pool_pre_ping,
1833
- SQLAlchemy may reuse dead pooled connections, causing SSL errors.
1834
-
1835
- pytest-xdist:
1836
- Each worker has its own isolated branch. Resets only affect that worker's
1837
- branch, so workers don't interfere with each other.
1838
-
1839
- Requires either:
1840
- - NEON_API_KEY and NEON_PROJECT_ID environment variables, or
1841
- - --neon-api-key and --neon-project-id command line options
1842
-
1843
- Yields:
1844
- NeonBranch: Object with branch_id, project_id, connection_string, and host.
1845
-
1846
- Example::
1847
-
1848
- def test_insert_user(neon_branch_isolated):
1849
- # DATABASE_URL is automatically set
1850
- conn_string = os.environ["DATABASE_URL"]
1851
-
1852
- # Insert data - branch will reset after this test
1853
- with psycopg.connect(conn_string) as conn:
1854
- conn.execute("INSERT INTO users (name) VALUES ('test')")
1855
- conn.commit()
1856
- # Next test starts with clean state
1857
- """
1858
- # DATABASE_URL is already set by _neon_isolated_branch
1859
- yield _neon_isolated_branch
1860
-
1861
- # Reset branch to migration state after each test
1862
- branch_manager = getattr(request.config, "_neon_isolated_branch_manager", None)
1863
- if branch_manager is not None:
1864
- try:
1865
- branch_manager.reset_branch(_neon_isolated_branch)
1866
- except Exception as e:
1867
- pytest.fail(
1868
- f"\n\nFailed to reset branch {_neon_isolated_branch.branch_id} "
1869
- f"after test. Subsequent tests may see dirty state.\n\n"
1870
- f"Error: {e}\n\n"
1871
- f"To keep the branch for debugging, use --neon-keep-branches"
1872
- )
1873
-
1874
-
1875
- @pytest.fixture(scope="function")
1876
- def neon_branch_readwrite(
1877
- neon_branch_isolated: NeonBranch,
1878
- ) -> Generator[NeonBranch, None, None]:
1879
- """
1880
- Deprecated: Use ``neon_branch_isolated`` instead.
1881
-
1882
- This fixture is now an alias for ``neon_branch_isolated``.
1883
-
1884
- .. deprecated:: 2.3.0
1885
- Use ``neon_branch_isolated`` for tests that modify data with reset,
1886
- ``neon_branch_dirty`` for shared state, or ``neon_branch_readonly``
1887
- for read-only access.
1888
- """
1889
- warnings.warn(
1890
- "neon_branch_readwrite is deprecated. Use neon_branch_isolated (for tests "
1891
- "that modify data with isolation) or neon_branch_dirty (for shared state).",
1892
- DeprecationWarning,
1893
- stacklevel=2,
1894
- )
1895
- yield neon_branch_isolated
1896
-
1897
-
1898
- @pytest.fixture(scope="function")
1899
- def neon_branch(
1900
- neon_branch_isolated: NeonBranch,
1901
- ) -> Generator[NeonBranch, None, None]:
1902
- """
1903
- Deprecated: Use ``neon_branch_isolated``, ``neon_branch_dirty``, or
1904
- ``neon_branch_readonly`` instead.
1905
-
1906
- This fixture is now an alias for ``neon_branch_isolated``.
1907
-
1908
- .. deprecated:: 1.1.0
1909
- Use ``neon_branch_isolated`` for tests that modify data with reset,
1910
- ``neon_branch_dirty`` for shared state, or ``neon_branch_readonly``
1911
- for read-only access.
1912
- """
1913
- warnings.warn(
1914
- "neon_branch is deprecated. Use neon_branch_isolated (for tests that "
1915
- "modify data), neon_branch_dirty (for shared state), or "
1916
- "neon_branch_readonly (for read-only tests).",
1917
- DeprecationWarning,
1918
- stacklevel=2,
1919
- )
1920
- yield neon_branch_isolated
1921
-
1922
-
1923
- @pytest.fixture(scope="module")
1924
- def neon_branch_shared(
1925
- request: pytest.FixtureRequest,
1926
- _neon_migration_branch: NeonBranch,
1927
- _neon_migrations_synchronized: Any, # Ensures migrations complete first
1928
- ) -> Generator[NeonBranch, None, None]:
1929
- """
1930
- Provide a shared Neon database branch for all tests in a module.
1931
-
1932
- This fixture creates one branch per test module and shares it across all
1933
- tests without resetting. This is the fastest option but tests can see
1934
- each other's data modifications.
1935
-
1936
- If you override the `neon_apply_migrations` fixture, migrations will run
1937
- once before the first test, and this branch will include the migrated schema.
1938
-
1939
- Use this when:
1940
- - Tests are read-only or don't interfere with each other
1941
- - You manually clean up test data within each test
1942
- - Maximum speed is more important than isolation
1943
-
1944
- Warning: Tests in the same module will share database state. Data created
1945
- by one test will be visible to subsequent tests. Use `neon_branch` instead
1946
- if you need isolation between tests.
1947
-
1948
- Yields:
1949
- 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
+ )
1950
1011
 
1951
- Example:
1952
- def test_read_only_query(neon_branch_shared):
1953
- # Fast: no reset between tests, but be careful about data leakage
1954
- conn_string = neon_branch_shared.connection_string
1955
- """
1956
- yield from _create_neon_branch(
1957
- request,
1958
- parent_branch_id_override=_neon_migration_branch.branch_id,
1959
- branch_name_suffix="-shared",
1960
- )
1012
+ return branch
1961
1013
 
1962
1014
 
1963
1015
  @pytest.fixture
1964
- def neon_connection(neon_branch_isolated: NeonBranch):
1016
+ def neon_connection(neon_branch: NeonBranch):
1965
1017
  """
1966
1018
  Provide a psycopg2 connection to the test branch.
1967
1019
 
@@ -1969,7 +1021,6 @@ def neon_connection(neon_branch_isolated: NeonBranch):
1969
1021
  pip install pytest-neon[psycopg2]
1970
1022
 
1971
1023
  The connection is rolled back and closed after each test.
1972
- Uses neon_branch_isolated for test isolation.
1973
1024
 
1974
1025
  Yields:
1975
1026
  psycopg2 connection object
@@ -1991,22 +1042,22 @@ def neon_connection(neon_branch_isolated: NeonBranch):
1991
1042
  " The 'neon_connection' fixture requires psycopg2.\n\n"
1992
1043
  " To fix this, install the psycopg2 extra:\n\n"
1993
1044
  " pip install pytest-neon[psycopg2]\n\n"
1994
- " Or use the 'neon_branch_isolated' fixture with your own driver:\n\n"
1995
- " 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"
1996
1047
  " import your_driver\n"
1997
1048
  " conn = your_driver.connect(\n"
1998
- " neon_branch_isolated.connection_string)\n\n"
1049
+ " neon_branch.connection_string)\n\n"
1999
1050
  "═══════════════════════════════════════════════════════════════════\n"
2000
1051
  )
2001
1052
 
2002
- conn = psycopg2.connect(neon_branch_isolated.connection_string)
1053
+ conn = psycopg2.connect(neon_branch.connection_string)
2003
1054
  yield conn
2004
1055
  conn.rollback()
2005
1056
  conn.close()
2006
1057
 
2007
1058
 
2008
1059
  @pytest.fixture
2009
- def neon_connection_psycopg(neon_branch_isolated: NeonBranch):
1060
+ def neon_connection_psycopg(neon_branch: NeonBranch):
2010
1061
  """
2011
1062
  Provide a psycopg (v3) connection to the test branch.
2012
1063
 
@@ -2014,7 +1065,6 @@ def neon_connection_psycopg(neon_branch_isolated: NeonBranch):
2014
1065
  pip install pytest-neon[psycopg]
2015
1066
 
2016
1067
  The connection is rolled back and closed after each test.
2017
- Uses neon_branch_isolated for test isolation.
2018
1068
 
2019
1069
  Yields:
2020
1070
  psycopg connection object
@@ -2036,40 +1086,29 @@ def neon_connection_psycopg(neon_branch_isolated: NeonBranch):
2036
1086
  " The 'neon_connection_psycopg' fixture requires psycopg v3.\n\n"
2037
1087
  " To fix this, install the psycopg extra:\n\n"
2038
1088
  " pip install pytest-neon[psycopg]\n\n"
2039
- " Or use the 'neon_branch_isolated' fixture with your own driver:\n\n"
2040
- " 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"
2041
1091
  " import your_driver\n"
2042
1092
  " conn = your_driver.connect(\n"
2043
- " neon_branch_isolated.connection_string)\n\n"
1093
+ " neon_branch.connection_string)\n\n"
2044
1094
  "═══════════════════════════════════════════════════════════════════\n"
2045
1095
  )
2046
1096
 
2047
- conn = psycopg.connect(neon_branch_isolated.connection_string)
1097
+ conn = psycopg.connect(neon_branch.connection_string)
2048
1098
  yield conn
2049
1099
  conn.rollback()
2050
1100
  conn.close()
2051
1101
 
2052
1102
 
2053
1103
  @pytest.fixture
2054
- def neon_engine(neon_branch_isolated: NeonBranch):
1104
+ def neon_engine(neon_branch: NeonBranch):
2055
1105
  """
2056
1106
  Provide a SQLAlchemy engine connected to the test branch.
2057
1107
 
2058
1108
  Requires the sqlalchemy optional dependency:
2059
1109
  pip install pytest-neon[sqlalchemy]
2060
1110
 
2061
- The engine is disposed after each test, which handles stale connections
2062
- after branch resets automatically. Uses neon_branch_isolated for test isolation.
2063
-
2064
- Note:
2065
- If you create your own module-level engine instead of using this
2066
- fixture, you MUST use pool_pre_ping=True::
2067
-
2068
- engine = create_engine(DATABASE_URL, pool_pre_ping=True)
2069
-
2070
- This is required because branch resets terminate server-side
2071
- connections, and without pool_pre_ping SQLAlchemy may reuse dead
2072
- pooled connections.
1111
+ The engine is disposed after each test.
2073
1112
 
2074
1113
  Yields:
2075
1114
  SQLAlchemy Engine object
@@ -2091,14 +1130,14 @@ def neon_engine(neon_branch_isolated: NeonBranch):
2091
1130
  " The 'neon_engine' fixture requires SQLAlchemy.\n\n"
2092
1131
  " To fix this, install the sqlalchemy extra:\n\n"
2093
1132
  " pip install pytest-neon[sqlalchemy]\n\n"
2094
- " Or use the 'neon_branch_isolated' fixture with your own driver:\n\n"
2095
- " 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"
2096
1135
  " from sqlalchemy import create_engine\n"
2097
1136
  " engine = create_engine(\n"
2098
- " neon_branch_isolated.connection_string)\n\n"
1137
+ " neon_branch.connection_string)\n\n"
2099
1138
  "═══════════════════════════════════════════════════════════════════\n"
2100
1139
  )
2101
1140
 
2102
- engine = create_engine(neon_branch_isolated.connection_string)
1141
+ engine = create_engine(neon_branch.connection_string)
2103
1142
  yield engine
2104
1143
  engine.dispose()