chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
Integration test suite for PostgreSQL retry mechanism using real database.
|
||||
|
||||
This test suite connects to a real PostgreSQL database using credentials from .env
|
||||
and tests the retry mechanism with actual network failures.
|
||||
|
||||
Prerequisites:
|
||||
1. PostgreSQL server running and accessible
|
||||
2. .env file with POSTGRES_* configuration
|
||||
3. asyncpg installed: pip install asyncpg
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from dotenv import load_dotenv
|
||||
from unittest.mock import patch
|
||||
from lightrag.kg.postgres_impl import PostgreSQLDB
|
||||
|
||||
asyncpg = pytest.importorskip("asyncpg")
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv(dotenv_path=".env", override=False)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_db
|
||||
class TestPostgresRetryIntegration:
|
||||
"""Integration tests for PostgreSQL retry mechanism with real database."""
|
||||
|
||||
@pytest.fixture
|
||||
def db_config(self):
|
||||
"""Load database configuration from environment variables.
|
||||
|
||||
Uses new HA-optimized defaults that match postgres_impl.py ClientManager.get_config():
|
||||
- 10 retry attempts (up from 3)
|
||||
- 3.0s initial backoff (up from 0.5s)
|
||||
- 30.0s max backoff (up from 5.0s)
|
||||
"""
|
||||
return {
|
||||
"host": os.getenv("POSTGRES_HOST", "localhost"),
|
||||
"port": int(os.getenv("POSTGRES_PORT", "5432")),
|
||||
"user": os.getenv("POSTGRES_USER", "postgres"),
|
||||
"password": os.getenv("POSTGRES_PASSWORD", ""),
|
||||
"database": os.getenv("POSTGRES_DATABASE", "postgres"),
|
||||
"workspace": os.getenv("POSTGRES_WORKSPACE", "test_retry"),
|
||||
"max_connections": int(os.getenv("POSTGRES_MAX_CONNECTIONS", "10")),
|
||||
# Connection retry configuration - mirrors postgres_impl.py ClientManager.get_config()
|
||||
# NEW DEFAULTS optimized for HA deployments
|
||||
"connection_retry_attempts": min(
|
||||
100,
|
||||
int(os.getenv("POSTGRES_CONNECTION_RETRIES", "10")), # 3 → 10
|
||||
),
|
||||
"connection_retry_backoff": min(
|
||||
300.0,
|
||||
float(
|
||||
os.getenv("POSTGRES_CONNECTION_RETRY_BACKOFF", "3.0")
|
||||
), # 0.5 → 3.0
|
||||
),
|
||||
"connection_retry_backoff_max": min(
|
||||
600.0,
|
||||
float(
|
||||
os.getenv("POSTGRES_CONNECTION_RETRY_BACKOFF_MAX", "30.0")
|
||||
), # 5.0 → 30.0
|
||||
),
|
||||
"pool_close_timeout": min(
|
||||
30.0, float(os.getenv("POSTGRES_POOL_CLOSE_TIMEOUT", "5.0"))
|
||||
),
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_connection_success(self, db_config):
|
||||
"""
|
||||
Test successful connection to real PostgreSQL database.
|
||||
|
||||
This validates that:
|
||||
1. Database credentials are correct
|
||||
2. Connection pool initializes properly
|
||||
3. Basic query works
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 1: Real Database Connection")
|
||||
print("=" * 80)
|
||||
print(
|
||||
f" → Connecting to {db_config['host']}:{db_config['port']}/{db_config['database']}"
|
||||
)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
# Initialize database connection
|
||||
await db.initdb()
|
||||
print(" ✓ Connection successful")
|
||||
|
||||
# Test simple query
|
||||
result = await db.query("SELECT 1 as test", multirows=False)
|
||||
assert result is not None
|
||||
assert result.get("test") == 1
|
||||
print(" ✓ Query executed successfully")
|
||||
|
||||
print("\n✅ Test passed: Real database connection works")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simulated_transient_error_with_real_db(self, db_config):
|
||||
"""
|
||||
Test retry mechanism with simulated transient errors on real database.
|
||||
|
||||
Simulates connection failures on first 2 attempts, then succeeds.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 2: Simulated Transient Errors")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
attempt_count = {"value": 0}
|
||||
|
||||
# Original create_pool function
|
||||
original_create_pool = asyncpg.create_pool
|
||||
|
||||
async def mock_create_pool_with_failures(*args, **kwargs):
|
||||
"""Mock that fails first 2 times, then calls real create_pool."""
|
||||
attempt_count["value"] += 1
|
||||
print(f" → Connection attempt {attempt_count['value']}")
|
||||
|
||||
if attempt_count["value"] <= 2:
|
||||
print(" ✗ Simulating connection failure")
|
||||
raise asyncpg.exceptions.ConnectionFailureError(
|
||||
f"Simulated failure on attempt {attempt_count['value']}"
|
||||
)
|
||||
|
||||
print(" ✓ Allowing real connection")
|
||||
return await original_create_pool(*args, **kwargs)
|
||||
|
||||
try:
|
||||
# Patch create_pool to simulate failures
|
||||
with patch(
|
||||
"asyncpg.create_pool", side_effect=mock_create_pool_with_failures
|
||||
):
|
||||
await db.initdb()
|
||||
|
||||
assert attempt_count["value"] == 3, (
|
||||
f"Expected 3 attempts, got {attempt_count['value']}"
|
||||
)
|
||||
assert db.pool is not None, "Pool should be initialized after retries"
|
||||
|
||||
# Verify database is actually working
|
||||
result = await db.query("SELECT 1 as test", multirows=False)
|
||||
assert result.get("test") == 1
|
||||
|
||||
print(
|
||||
f"\n✅ Test passed: Retry mechanism worked, connected after {attempt_count['value']} attempts"
|
||||
)
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_retry_with_real_db(self, db_config):
|
||||
"""
|
||||
Test query-level retry with simulated connection issues.
|
||||
|
||||
Tests that queries retry on transient failures by simulating
|
||||
a temporary database unavailability.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 3: Query-Level Retry")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
# First initialize normally
|
||||
await db.initdb()
|
||||
print(" ✓ Database initialized")
|
||||
|
||||
# Close the pool to simulate connection loss
|
||||
print(" → Simulating connection loss (closing pool)...")
|
||||
await db.pool.close()
|
||||
db.pool = None
|
||||
|
||||
# Now query should trigger pool recreation and retry
|
||||
print(" → Attempting query (should auto-reconnect)...")
|
||||
result = await db.query("SELECT 1 as test", multirows=False)
|
||||
|
||||
assert result.get("test") == 1, "Query should succeed after reconnection"
|
||||
assert db.pool is not None, "Pool should be recreated"
|
||||
|
||||
print(" ✓ Query succeeded after automatic reconnection")
|
||||
print("\n✅ Test passed: Auto-reconnection works correctly")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_queries_with_real_db(self, db_config):
|
||||
"""
|
||||
Test concurrent queries to validate thread safety and connection pooling.
|
||||
|
||||
Runs multiple concurrent queries to ensure no deadlocks or race conditions.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 4: Concurrent Queries")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
await db.initdb()
|
||||
print(" ✓ Database initialized")
|
||||
|
||||
# Launch 10 concurrent queries
|
||||
num_queries = 10
|
||||
print(f" → Launching {num_queries} concurrent queries...")
|
||||
|
||||
async def run_query(query_id):
|
||||
result = await db.query(
|
||||
f"SELECT {query_id} as id, pg_sleep(0.1)", multirows=False
|
||||
)
|
||||
return result.get("id")
|
||||
|
||||
start_time = time.time()
|
||||
tasks = [run_query(i) for i in range(num_queries)]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Check results
|
||||
successful = sum(1 for r in results if not isinstance(r, Exception))
|
||||
failed = sum(1 for r in results if isinstance(r, Exception))
|
||||
|
||||
print(f" → Completed in {elapsed:.2f}s")
|
||||
print(f" → Results: {successful} successful, {failed} failed")
|
||||
|
||||
assert successful == num_queries, (
|
||||
f"All {num_queries} queries should succeed"
|
||||
)
|
||||
assert failed == 0, "No queries should fail"
|
||||
|
||||
print("\n✅ Test passed: All concurrent queries succeeded, no deadlocks")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_close_timeout_real(self, db_config):
|
||||
"""
|
||||
Test pool close timeout protection with real database.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 5: Pool Close Timeout")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
await db.initdb()
|
||||
print(" ✓ Database initialized")
|
||||
|
||||
# Trigger pool reset (which includes close)
|
||||
print(" → Triggering pool reset...")
|
||||
start_time = time.time()
|
||||
await db._reset_pool()
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
print(f" ✓ Pool reset completed in {elapsed:.2f}s")
|
||||
assert db.pool is None, "Pool should be None after reset"
|
||||
assert elapsed < db.pool_close_timeout + 1, (
|
||||
"Reset should complete within timeout"
|
||||
)
|
||||
|
||||
print("\n✅ Test passed: Pool reset handled correctly")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
# Already closed in test
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_configuration_from_env(self, db_config):
|
||||
"""
|
||||
Test that configuration is correctly loaded from environment variables.
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 6: Environment Configuration")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
print(" → Configuration loaded:")
|
||||
print(f" • Host: {db.host}")
|
||||
print(f" • Port: {db.port}")
|
||||
print(f" • Database: {db.database}")
|
||||
print(f" • User: {db.user}")
|
||||
print(f" • Workspace: {db.workspace}")
|
||||
print(f" • Max Connections: {db.max}")
|
||||
print(f" • Retry Attempts: {db.connection_retry_attempts}")
|
||||
print(f" • Retry Backoff: {db.connection_retry_backoff}s")
|
||||
print(f" • Max Backoff: {db.connection_retry_backoff_max}s")
|
||||
print(f" • Pool Close Timeout: {db.pool_close_timeout}s")
|
||||
|
||||
# Verify required fields are present
|
||||
assert db.host, "Host should be configured"
|
||||
assert db.port, "Port should be configured"
|
||||
assert db.user, "User should be configured"
|
||||
assert db.database, "Database should be configured"
|
||||
|
||||
print("\n✅ Test passed: All configuration loaded correctly from .env")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
def run_integration_tests():
|
||||
"""Run all integration tests with detailed output."""
|
||||
print("\n" + "=" * 80)
|
||||
print("POSTGRESQL RETRY MECHANISM - INTEGRATION TESTS")
|
||||
print("Testing with REAL database from .env configuration")
|
||||
print("=" * 80)
|
||||
|
||||
# Check if database configuration exists
|
||||
if not os.getenv("POSTGRES_HOST"):
|
||||
print("\n⚠️ WARNING: No POSTGRES_HOST in .env file")
|
||||
print("Please ensure .env file exists with PostgreSQL configuration.")
|
||||
return
|
||||
|
||||
print("\nRunning integration tests...\n")
|
||||
|
||||
# Run pytest with verbose output
|
||||
pytest.main(
|
||||
[
|
||||
__file__,
|
||||
"-v",
|
||||
"-s", # Don't capture output
|
||||
"--tb=short", # Short traceback format
|
||||
"--color=yes",
|
||||
"-x", # Stop on first failure
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_integration_tests()
|
||||
Reference in New Issue
Block a user