77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
"""Tests for shared utilities."""
|
|
|
|
import logging
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from llm_inference.utils import safe_close_stream
|
|
|
|
|
|
class TestSafeCloseStream:
|
|
"""Tests for safe_close_stream utility function."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_closes_stream_with_aclose(self) -> None:
|
|
"""Test closing stream that has aclose() method."""
|
|
stream = AsyncMock()
|
|
stream.aclose = AsyncMock()
|
|
logger = MagicMock(spec=logging.Logger)
|
|
|
|
await safe_close_stream(stream, logger)
|
|
|
|
stream.aclose.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_closes_stream_with_close_fallback(self) -> None:
|
|
"""Test closing stream that only has close() method."""
|
|
stream = AsyncMock()
|
|
del stream.aclose # Remove aclose to test fallback
|
|
stream.close = AsyncMock()
|
|
logger = MagicMock(spec=logging.Logger)
|
|
|
|
await safe_close_stream(stream, logger)
|
|
|
|
stream.close.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handles_none_stream(self) -> None:
|
|
"""Test that None stream is handled gracefully."""
|
|
logger = MagicMock(spec=logging.Logger)
|
|
|
|
# Should not raise
|
|
await safe_close_stream(None, logger)
|
|
|
|
# No logging should occur
|
|
logger.debug.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handles_close_exception(self) -> None:
|
|
"""Test that exceptions during close are caught and logged."""
|
|
stream = AsyncMock()
|
|
stream.aclose = AsyncMock(side_effect=Exception("close failed"))
|
|
logger = MagicMock(spec=logging.Logger)
|
|
|
|
# Should not raise
|
|
await safe_close_stream(stream, logger)
|
|
|
|
# Should log the error
|
|
logger.debug.assert_called_once()
|
|
assert "close failed" in str(logger.debug.call_args)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handles_stream_without_close_methods(self) -> None:
|
|
"""Test handling stream with no close methods."""
|
|
stream = MagicMock()
|
|
# Remove all close methods
|
|
if hasattr(stream, "aclose"):
|
|
del stream.aclose
|
|
if hasattr(stream, "close"):
|
|
del stream.close
|
|
logger = MagicMock(spec=logging.Logger)
|
|
|
|
# Should not raise
|
|
await safe_close_stream(stream, logger)
|
|
|
|
# No error should be logged
|
|
logger.debug.assert_not_called()
|