| 1 | import pytest |
| 2 | |
| 3 | from helpers import extension |
| 4 | from helpers.llm_result import LLMResult |
| 5 | from models import LiteLLMChatWrapper |
| 6 | |
| 7 | |
| 8 | @pytest.mark.asyncio |
| 9 | @pytest.mark.parametrize( |
| 10 | ("method_name", "result"), |
| 11 | [ |
| 12 | ("unified_call", ("response", "reasoning")), |
| 13 | ("unified_turn", LLMResult(response="response")), |
| 14 | ], |
| 15 | ) |
| 16 | async def test_unified_model_calls_expose_function_extensions( |
| 17 | monkeypatch, method_name, result |
| 18 | ): |
| 19 | points = [] |
| 20 | |
| 21 | async def call_extensions(point, agent=None, **kwargs): |
| 22 | points.append(point) |
| 23 | if point.endswith("/start"): |
| 24 | kwargs["data"]["result"] = result |
| 25 | |
| 26 | monkeypatch.setattr(extension, "call_extensions_async", call_extensions) |
| 27 | |
| 28 | actual = await getattr(LiteLLMChatWrapper, method_name)(object()) |
| 29 | |
| 30 | assert actual is result |
| 31 | assert points == [ |
| 32 | f"_functions/models/LiteLLMChatWrapper/{method_name}/start", |
| 33 | f"_functions/models/LiteLLMChatWrapper/{method_name}/end", |
| 34 | ] |