correct type hints for mocks

huggingface · Feb 13, 2025 · 893e409 · 893e409
1 parent 02faf60
commit 893e409
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/tests/test_agents.py b/tests/test_agents.py
@@ -143,7 +143,7 @@ def __call__(self, messages, tools_to_call_from=None, stop_sequences=None, gramm
             )
 
 
-def fake_code_model(messages, stop_sequences=None, grammar=None) -> str:
+def fake_code_model(messages, stop_sequences=None, grammar=None) -> ChatMessage:
     prompt = str(messages)
     if "special_marker" not in prompt:
         return ChatMessage(
@@ -169,7 +169,7 @@ def fake_code_model(messages, stop_sequences=None, grammar=None) -> str:
         )
 
 
-def fake_code_model_error(messages, stop_sequences=None) -> str:
+def fake_code_model_error(messages, stop_sequences=None) -> ChatMessage:
     prompt = str(messages)
     if "special_marker" not in prompt:
         return ChatMessage(
@@ -199,7 +199,7 @@ def error_function():
         )
 
 
-def fake_code_model_syntax_error(messages, stop_sequences=None) -> str:
+def fake_code_model_syntax_error(messages, stop_sequences=None) -> ChatMessage:
     prompt = str(messages)
     if "special_marker" not in prompt:
         return ChatMessage(
@@ -228,7 +228,7 @@ def fake_code_model_syntax_error(messages, stop_sequences=None) -> str:
         )
 
 
-def fake_code_model_import(messages, stop_sequences=None) -> str:
+def fake_code_model_import(messages, stop_sequences=None) -> ChatMessage:
     return ChatMessage(
         role="assistant",
         content="""
@@ -242,7 +242,7 @@ def fake_code_model_import(messages, stop_sequences=None) -> str:
     )
 
 
-def fake_code_functiondef(messages, stop_sequences=None) -> str:
+def fake_code_functiondef(messages, stop_sequences=None) -> ChatMessage:
     prompt = str(messages)
     if "special_marker" not in prompt:
         return ChatMessage(
@@ -273,7 +273,7 @@ def moving_average(x, w):
         )
 
 
-def fake_code_model_single_step(messages, stop_sequences=None, grammar=None) -> str:
+def fake_code_model_single_step(messages, stop_sequences=None, grammar=None) -> ChatMessage:
     return ChatMessage(
         role="assistant",
         content="""
@@ -287,7 +287,7 @@ def fake_code_model_single_step(messages, stop_sequences=None, grammar=None) ->
     )
 
 
-def fake_code_model_no_return(messages, stop_sequences=None, grammar=None) -> str:
+def fake_code_model_no_return(messages, stop_sequences=None, grammar=None) -> ChatMessage:
     return ChatMessage(
         role="assistant",
         content="""