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/__init__.py +3 -3
- pytest_neon/plugin.py +140 -1084
- pytest_neon-3.0.0.dist-info/METADATA +348 -0
- pytest_neon-3.0.0.dist-info/RECORD +8 -0
- pytest_neon-2.3.1.dist-info/METADATA +0 -650
- pytest_neon-2.3.1.dist-info/RECORD +0 -8
- {pytest_neon-2.3.1.dist-info → pytest_neon-3.0.0.dist-info}/WHEEL +0 -0
- {pytest_neon-2.3.1.dist-info → pytest_neon-3.0.0.dist-info}/entry_points.txt +0 -0
- {pytest_neon-2.3.1.dist-info → pytest_neon-3.0.0.dist-info}/licenses/LICENSE +0 -0
pytest_neon/plugin.py
CHANGED
|
@@ -1,15 +1,10 @@
|
|
|
1
1
|
"""Pytest plugin providing Neon database branch fixtures.
|
|
2
2
|
|
|
3
|
-
This plugin provides
|
|
4
|
-
branching feature.
|
|
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
|
|
7
|
-
|
|
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
|
-
└──
|
|
22
|
-
|
|
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
|
|
249
|
-
|
|
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
|
|
299
|
-
|
|
300
|
-
|
|
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
|
|
297
|
+
Get the password for a role WITHOUT resetting it.
|
|
312
298
|
|
|
313
|
-
|
|
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
|
-
|
|
318
|
-
|
|
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
|
-
|
|
322
|
-
except
|
|
323
|
-
|
|
324
|
-
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
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., "-
|
|
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.
|
|
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
|
|
549
|
+
def _get_password_and_build_connection_string(
|
|
636
550
|
self, branch_id: str, host: str
|
|
637
551
|
) -> str:
|
|
638
|
-
"""
|
|
639
|
-
|
|
640
|
-
lambda:
|
|
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="
|
|
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
|
|
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
|
|
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
|
|
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
|
-
|
|
868
|
+
Internal: Create test branch, coordinated across workers.
|
|
1412
869
|
|
|
1413
|
-
This
|
|
1414
|
-
|
|
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
|
-
|
|
1418
|
-
|
|
1419
|
-
|
|
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
|
|
879
|
+
def create_branch() -> dict[str, Any]:
|
|
1433
880
|
b = _neon_branch_manager.create_branch(
|
|
1434
|
-
name_suffix="-
|
|
1435
|
-
expiry_seconds=
|
|
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
|
-
"
|
|
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(
|
|
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
|
|
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
|
|
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(
|
|
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(
|
|
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(
|
|
931
|
+
def neon_apply_migrations(_neon_test_branch):
|
|
1504
932
|
import psycopg
|
|
1505
|
-
|
|
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
|
-
|
|
1512
|
-
Use
|
|
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
|
|
1517
|
-
the fixture was not overridden.
|
|
945
|
+
Any value (ignored). The default returns None.
|
|
1518
946
|
"""
|
|
1519
|
-
return
|
|
947
|
+
return None
|
|
1520
948
|
|
|
1521
949
|
|
|
1522
950
|
@pytest.fixture(scope="session")
|
|
1523
|
-
def
|
|
1524
|
-
|
|
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
|
|
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
|
|
1694
|
-
|
|
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
|
-
|
|
1697
|
-
|
|
1698
|
-
|
|
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(
|
|
979
|
+
def test_query_users(neon_branch):
|
|
1713
980
|
# DATABASE_URL is automatically set
|
|
1714
|
-
|
|
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)
|
|
984
|
+
assert len(result) >= 0
|
|
1720
985
|
|
|
1721
|
-
|
|
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
|
-
|
|
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
|
-
|
|
1773
|
-
|
|
1774
|
-
|
|
1775
|
-
|
|
1776
|
-
|
|
1777
|
-
|
|
1778
|
-
|
|
1779
|
-
|
|
1780
|
-
|
|
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
|
-
|
|
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
|
-
|
|
1814
|
-
|
|
1815
|
-
|
|
1816
|
-
|
|
1817
|
-
|
|
1818
|
-
|
|
1819
|
-
|
|
1820
|
-
|
|
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
|
-
|
|
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(
|
|
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 '
|
|
1978
|
-
" def test_example(
|
|
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
|
-
"
|
|
1049
|
+
" neon_branch.connection_string)\n\n"
|
|
1982
1050
|
"═══════════════════════════════════════════════════════════════════\n"
|
|
1983
1051
|
)
|
|
1984
1052
|
|
|
1985
|
-
conn = psycopg2.connect(
|
|
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(
|
|
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 '
|
|
2023
|
-
" def test_example(
|
|
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
|
-
"
|
|
1093
|
+
" neon_branch.connection_string)\n\n"
|
|
2027
1094
|
"═══════════════════════════════════════════════════════════════════\n"
|
|
2028
1095
|
)
|
|
2029
1096
|
|
|
2030
|
-
conn = psycopg.connect(
|
|
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(
|
|
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
|
|
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 '
|
|
2078
|
-
" def test_example(
|
|
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
|
-
"
|
|
1137
|
+
" neon_branch.connection_string)\n\n"
|
|
2082
1138
|
"═══════════════════════════════════════════════════════════════════\n"
|
|
2083
1139
|
)
|
|
2084
1140
|
|
|
2085
|
-
engine = create_engine(
|
|
1141
|
+
engine = create_engine(neon_branch.connection_string)
|
|
2086
1142
|
yield engine
|
|
2087
1143
|
engine.dispose()
|