Skip to content

Commit 538a5b0

Browse files
aperepellizzij
authored andcommitted
refactor: use _FINISH_REASON_MAPPING directly in tests
Import and use the actual _FINISH_REASON_MAPPING from lite_llm instead of duplicating it in tests. This ensures tests stay in sync with implementation changes automatically. Benefits: - Single source of truth for finish_reason mappings - Tests automatically reflect any future mapping changes - Reduced code duplication - Better maintainability Addresses review comment: google#3114 (review)
1 parent 692af95 commit 538a5b0

1 file changed

Lines changed: 3 additions & 8 deletions

File tree

tests/unittests/models/test_litellm.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import warnings
2020

2121
from google.adk.models.lite_llm import _content_to_message_param
22+
from google.adk.models.lite_llm import _FINISH_REASON_MAPPING
2223
from google.adk.models.lite_llm import _function_declaration_to_tool_param
2324
from google.adk.models.lite_llm import _get_content
2425
from google.adk.models.lite_llm import _message_to_generate_content_response
@@ -1969,14 +1970,8 @@ async def test_finish_reason_propagation(
19691970
assert response.content.role == "model"
19701971
# Verify finish_reason is mapped to FinishReason enum
19711972
assert isinstance(response.finish_reason, types.FinishReason)
1972-
# Verify correct enum mapping using dictionary
1973-
expected_mapping = {
1974-
"length": types.FinishReason.MAX_TOKENS,
1975-
"stop": types.FinishReason.STOP,
1976-
"tool_calls": types.FinishReason.STOP,
1977-
"content_filter": types.FinishReason.SAFETY,
1978-
}
1979-
assert response.finish_reason == expected_mapping[finish_reason]
1973+
# Verify correct enum mapping using the actual mapping from lite_llm
1974+
assert response.finish_reason == _FINISH_REASON_MAPPING[finish_reason]
19801975
if expected_content:
19811976
assert response.content.parts[0].text == expected_content
19821977
if has_tool_calls:

0 commit comments

Comments
 (0)