Livrare LOT 1 - Didi

ai_platform/modules/llm-inference/tests/test_utils.py

@@ -0,0 +1,77 @@
+"""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()