66 lines
1.5 KiB
Python
66 lines
1.5 KiB
Python
|
"""Tests for connection pool."""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from fusionagi.api.pool import ConnectionPool
|
||
|
|
|
||
|
|
|
||
|
|
class MockConnection:
|
||
|
|
"""Mock connection for testing."""
|
||
|
|
def __init__(self):
|
||
|
|
self.connected = False
|
||
|
|
self.closed = False
|
||
|
|
|
||
|
|
async def connect(self):
|
||
|
|
self.connected = True
|
||
|
|
|
||
|
|
async def close(self):
|
||
|
|
self.closed = True
|
||
|
|
|
||
|
|
def is_alive(self):
|
||
|
|
return self.connected and not self.closed
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def pool():
|
||
|
|
return ConnectionPool(factory=MockConnection, min_size=2, max_size=5)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_initialize(pool):
|
||
|
|
await pool.initialize()
|
||
|
|
stats = pool.stats()
|
||
|
|
assert stats["available"] == 2
|
||
|
|
assert stats["total_created"] == 2
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_acquire_and_release(pool):
|
||
|
|
await pool.initialize()
|
||
|
|
conn = await pool.acquire()
|
||
|
|
assert isinstance(conn, MockConnection)
|
||
|
|
stats = pool.stats()
|
||
|
|
assert stats["in_use"] == 1
|
||
|
|
await pool.release(conn)
|
||
|
|
stats = pool.stats()
|
||
|
|
assert stats["in_use"] == 0
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_close_all(pool):
|
||
|
|
await pool.initialize()
|
||
|
|
await pool.close_all()
|
||
|
|
stats = pool.stats()
|
||
|
|
assert stats["available"] == 0
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_max_size():
|
||
|
|
pool = ConnectionPool(factory=MockConnection, min_size=1, max_size=2)
|
||
|
|
await pool.initialize()
|
||
|
|
c1 = await pool.acquire()
|
||
|
|
c2 = await pool.acquire()
|
||
|
|
assert pool.stats()["in_use"] == 2
|
||
|
|
await pool.release(c1)
|
||
|
|
await pool.release(c2)
|