BioMCP

"""Tests for connection pool management."""

import asyncio
import ssl
import weakref
from unittest.mock import patch

import httpx
import pytest

from biomcp.connection_pool import (
    EventLoopConnectionPools,
    close_all_pools,
    get_connection_pool,
)


@pytest.fixture
def pool_manager():
    """Create a fresh pool manager for testing."""
    return EventLoopConnectionPools()


@pytest.mark.asyncio
async def test_get_pool_creates_new_pool(pool_manager):
    """Test that get_pool creates a new pool when none exists."""
    timeout = httpx.Timeout(30)

    pool = await pool_manager.get_pool(verify=True, timeout=timeout)

    assert pool is not None
    assert isinstance(pool, httpx.AsyncClient)
    assert not pool.is_closed


@pytest.mark.asyncio
async def test_get_pool_reuses_existing_pool(pool_manager):
    """Test that get_pool reuses existing pools."""
    timeout = httpx.Timeout(30)

    pool1 = await pool_manager.get_pool(verify=True, timeout=timeout)
    pool2 = await pool_manager.get_pool(verify=True, timeout=timeout)

    assert pool1 is pool2


@pytest.mark.asyncio
async def test_get_pool_different_verify_settings(pool_manager):
    """Test that different verify settings create different pools."""
    timeout = httpx.Timeout(30)

    pool1 = await pool_manager.get_pool(verify=True, timeout=timeout)
    pool2 = await pool_manager.get_pool(verify=False, timeout=timeout)

    assert pool1 is not pool2


@pytest.mark.asyncio
async def test_get_pool_ssl_context(pool_manager):
    """Test pool creation with SSL context."""
    ssl_context = ssl.create_default_context()
    timeout = httpx.Timeout(30)

    pool = await pool_manager.get_pool(verify=ssl_context, timeout=timeout)

    assert pool is not None
    assert isinstance(pool, httpx.AsyncClient)


@pytest.mark.asyncio
async def test_pool_cleanup_on_close_all(pool_manager):
    """Test that close_all properly closes all pools."""
    timeout = httpx.Timeout(30)

    await pool_manager.get_pool(verify=True, timeout=timeout)
    await pool_manager.get_pool(verify=False, timeout=timeout)

    await pool_manager.close_all()

    # After close_all, pools should be cleared
    assert len(pool_manager._loop_pools) == 0


@pytest.mark.asyncio
async def test_no_event_loop_returns_single_use_client(pool_manager):
    """Test behavior when no event loop is running."""
    with patch("asyncio.get_running_loop", side_effect=RuntimeError):
        timeout = httpx.Timeout(30)

        pool = await pool_manager.get_pool(verify=True, timeout=timeout)

        assert pool is not None
        # Single-use client should have no keepalive
        # Note: httpx client internal structure may vary


@pytest.mark.asyncio
async def test_pool_recreation_after_close(pool_manager):
    """Test that a new pool is created after the old one is closed."""
    timeout = httpx.Timeout(30)

    pool1 = await pool_manager.get_pool(verify=True, timeout=timeout)
    await pool1.aclose()

    pool2 = await pool_manager.get_pool(verify=True, timeout=timeout)

    assert pool1 is not pool2
    assert pool1.is_closed
    assert not pool2.is_closed


@pytest.mark.asyncio
async def test_weak_reference_cleanup():
    """Test that weak references are used for event loops."""
    pool_manager = EventLoopConnectionPools()

    # Verify that the pool manager uses weak references
    assert isinstance(pool_manager._loop_pools, weakref.WeakKeyDictionary)

    # Create a pool
    timeout = httpx.Timeout(30)
    pool = await pool_manager.get_pool(verify=True, timeout=timeout)

    # Verify pool was created
    assert pool is not None

    # The current event loop should be in the weak key dict
    current_loop = asyncio.get_running_loop()
    assert current_loop in pool_manager._loop_pools


@pytest.mark.asyncio
async def test_global_get_connection_pool():
    """Test the global get_connection_pool function."""
    with patch.dict("os.environ", {"BIOMCP_USE_CONNECTION_POOL": "true"}):
        timeout = httpx.Timeout(30)

        pool = await get_connection_pool(verify=True, timeout=timeout)

        assert pool is not None
        assert isinstance(pool, httpx.AsyncClient)


@pytest.mark.asyncio
async def test_global_close_all_pools():
    """Test the global close_all_pools function."""
    # Create some pools
    timeout = httpx.Timeout(30)
    await get_connection_pool(verify=True, timeout=timeout)
    await get_connection_pool(verify=False, timeout=timeout)

    # Close all pools
    await close_all_pools()

    # Verify cleanup (this is implementation-specific)
    from biomcp.connection_pool import _pool_manager

    assert len(_pool_manager._loop_pools) == 0


@pytest.mark.asyncio
async def test_concurrent_pool_creation(pool_manager):
    """Test thread-safe pool creation under concurrent access."""
    timeout = httpx.Timeout(30)

    async def get_pool():
        return await pool_manager.get_pool(verify=True, timeout=timeout)

    # Create 10 concurrent requests for the same pool
    pools = await asyncio.gather(*[get_pool() for _ in range(10)])

    # All should return the same pool instance
    assert all(pool is pools[0] for pool in pools)


@pytest.mark.asyncio
async def test_connection_pool_limits():
    """Test that connection pools have proper limits set."""
    pool_manager = EventLoopConnectionPools()
    timeout = httpx.Timeout(30)

    pool = await pool_manager.get_pool(verify=True, timeout=timeout)

    # Verify pool was created (actual limits are internal to httpx)
    assert pool is not None
    assert isinstance(pool, httpx.AsyncClient)