|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the Licens |
14 | 14 |
|
| 15 | +import contextlib |
15 | 16 | import json |
| 17 | +import logging |
| 18 | +import os |
| 19 | +import sys |
| 20 | +import tempfile |
| 21 | +import unittest |
16 | 22 | from unittest.mock import ANY |
17 | 23 | from unittest.mock import AsyncMock |
18 | 24 | from unittest.mock import Mock |
|
29 | 35 | from google.adk.models.lite_llm import _model_response_to_chunk |
30 | 36 | from google.adk.models.lite_llm import _model_response_to_generate_content_response |
31 | 37 | from google.adk.models.lite_llm import _parse_tool_calls_from_text |
| 38 | +from google.adk.models.lite_llm import _redirect_litellm_loggers_to_stdout |
32 | 39 | from google.adk.models.lite_llm import _schema_to_dict |
33 | 40 | from google.adk.models.lite_llm import _split_message_content_and_tool_calls |
34 | 41 | from google.adk.models.lite_llm import _to_litellm_response_format |
@@ -3133,3 +3140,79 @@ async def test_get_completion_inputs_non_openai_no_file_upload(mocker): |
3133 | 3140 | assert "file_id" not in content[1]["file"] |
3134 | 3141 |
|
3135 | 3142 | mock_acreate_file.assert_not_called() |
| 3143 | + |
| 3144 | + |
| 3145 | +class TestRedirectLitellmLoggersToStdout(unittest.TestCase): |
| 3146 | + """Tests for _redirect_litellm_loggers_to_stdout function.""" |
| 3147 | + |
| 3148 | + def test_redirects_stderr_handler_to_stdout(self): |
| 3149 | + """Test that handlers pointing to stderr are redirected to stdout.""" |
| 3150 | + test_logger = logging.getLogger("LiteLLM") |
| 3151 | + # Create a handler pointing to stderr |
| 3152 | + handler = logging.StreamHandler(sys.stderr) |
| 3153 | + test_logger.addHandler(handler) |
| 3154 | + |
| 3155 | + try: |
| 3156 | + self.assertIs(handler.stream, sys.stderr) |
| 3157 | + |
| 3158 | + _redirect_litellm_loggers_to_stdout() |
| 3159 | + |
| 3160 | + self.assertIs(handler.stream, sys.stdout) |
| 3161 | + finally: |
| 3162 | + # Clean up |
| 3163 | + test_logger.removeHandler(handler) |
| 3164 | + |
| 3165 | + def test_preserves_stdout_handler(self): |
| 3166 | + """Test that handlers already pointing to stdout are not modified.""" |
| 3167 | + test_logger = logging.getLogger("LiteLLM Proxy") |
| 3168 | + # Create a handler already pointing to stdout |
| 3169 | + handler = logging.StreamHandler(sys.stdout) |
| 3170 | + test_logger.addHandler(handler) |
| 3171 | + |
| 3172 | + try: |
| 3173 | + _redirect_litellm_loggers_to_stdout() |
| 3174 | + |
| 3175 | + self.assertIs(handler.stream, sys.stdout) |
| 3176 | + finally: |
| 3177 | + # Clean up |
| 3178 | + test_logger.removeHandler(handler) |
| 3179 | + |
| 3180 | + def test_does_not_affect_non_stream_handlers(self): |
| 3181 | + """Test that non-StreamHandler handlers are not affected.""" |
| 3182 | + test_logger = logging.getLogger("LiteLLM Router") |
| 3183 | + # Create a FileHandler (not a StreamHandler) |
| 3184 | + with tempfile.NamedTemporaryFile(delete=False) as temp_file: |
| 3185 | + temp_file_name = temp_file.name |
| 3186 | + with contextlib.closing( |
| 3187 | + logging.FileHandler(temp_file_name) |
| 3188 | + ) as file_handler: |
| 3189 | + test_logger.addHandler(file_handler) |
| 3190 | + |
| 3191 | + try: |
| 3192 | + _redirect_litellm_loggers_to_stdout() |
| 3193 | + # FileHandler should not be modified (it doesn't point to stderr or stdout) |
| 3194 | + self.assertEqual(file_handler.baseFilename, temp_file_name) |
| 3195 | + finally: |
| 3196 | + # Clean up |
| 3197 | + test_logger.removeHandler(file_handler) |
| 3198 | + os.unlink(temp_file_name) |
| 3199 | + |
| 3200 | + |
| 3201 | +@pytest.mark.parametrize( |
| 3202 | + "logger_name", |
| 3203 | + ["LiteLLM", "LiteLLM Proxy", "LiteLLM Router"], |
| 3204 | + ids=["LiteLLM", "LiteLLM Proxy", "LiteLLM Router"], |
| 3205 | +) |
| 3206 | +def test_handles_litellm_logger_names(logger_name): |
| 3207 | + """Test that LiteLLM logger names are processed.""" |
| 3208 | + test_logger = logging.getLogger(logger_name) |
| 3209 | + handler = logging.StreamHandler(sys.stderr) |
| 3210 | + test_logger.addHandler(handler) |
| 3211 | + |
| 3212 | + try: |
| 3213 | + _redirect_litellm_loggers_to_stdout() |
| 3214 | + |
| 3215 | + assert handler.stream is sys.stdout |
| 3216 | + finally: |
| 3217 | + # Clean up |
| 3218 | + test_logger.removeHandler(handler) |
0 commit comments