phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-24 18:24:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add router_obj tests.

Browse source
This commit is contained in:
Jun Komoda 2025-04-12 00:29:58 +00:00
parent 2498092343
commit c63859a136
1 changed files with 29 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

29
tests/litellm/test_main.py
View file

@ -390,3 +390,32 @@ class Test_Chat:
call(model="gemini/gemini-1.5-flash", messages=[{"role": "user", "content": "hello"}], bar="foo"),
]
assert params == {"acompletion": True}
def test_calls_completion_with_router_obj(self, mocker):
router_obj = mocker.MagicMock()
chatobj = litellm.main.Chat({}, router_obj=router_obj)
chatobj.completions.create(
messages=[{"role": "user", "content": "hello"}],
model="gemini/gemini-1.5-flash",
foo="bar",
)
router_obj.completion.assert_called_once_with(
model="gemini/gemini-1.5-flash",
messages=[{"role": "user", "content": "hello"}],
foo="bar"
)
@pytest.mark.asyncio
async def test_calls_acompletion_with_router_obj(self, mocker):
router_obj = mocker.AsyncMock()
chatobj = litellm.main.Chat({"acompletion": True}, router_obj=router_obj)
await chatobj.completions.create( # type: ignore
messages=[{"role": "user", "content": "hello"}],
model="gemini/gemini-1.5-flash",
foo="bar",
)
router_obj.acompletion.assert_called_once_with(
model="gemini/gemini-1.5-flash",
messages=[{"role": "user", "content": "hello"}],
foo="bar"
)