test(code_execution): add tests for _maybe_recover_from_api_rejection

Author: JesserHamdaoui

tests/unittests/flows/llm_flows/test_code_execution.py

@@ -24,6 +24,7 @@
 from google.adk.code_executors.built_in_code_executor import BuiltInCodeExecutor
 from google.adk.code_executors.code_execution_utils import CodeExecutionResult
 from google.adk.flows.llm_flows._code_execution import _extract_code_from_error_message
+from google.adk.flows.llm_flows._code_execution import _maybe_recover_from_api_rejection
 from google.adk.flows.llm_flows._code_execution import response_processor
 from google.adk.models.llm_response import LlmResponse
 from google.genai import types
@@ -174,3 +175,55 @@ async def test_logs_executed_code(mock_logger):
   mock_logger.debug.assert_called_once_with(
       'Executed code:\n```\n%s\n```', 'print("hello")'
   )
+
+
+# ---------------------------------------------------------------------------
+# _maybe_recover_from_api_rejection
+# ---------------------------------------------------------------------------
+
+
+def _make_rejected_response(error_code: str, code_snippet: str) -> LlmResponse:
+  return LlmResponse(
+      content=None,
+      error_code=error_code,
+      error_message=f'Unexpected tool call: {code_snippet}',
+  )
+
+
+def test_maybe_recover_unexpected_tool_call():
+  llm_response = _make_rejected_response('UNEXPECTED_TOOL_CALL', 'print(42)')
+  recovered = _maybe_recover_from_api_rejection(llm_response)
+
+  assert recovered is True
+  assert llm_response.content is not None
+  assert len(llm_response.content.parts) == 1
+  assert llm_response.content.parts[0].executable_code.code == 'print(42)'
+  assert llm_response.error_code is None
+  assert llm_response.error_message is None
+  assert llm_response.finish_reason is None
+
+
+def test_maybe_recover_malformed_function_call():
+  llm_response = _make_rejected_response('MALFORMED_FUNCTION_CALL', 'x=1')
+  assert _maybe_recover_from_api_rejection(llm_response) is True
+  assert llm_response.content is not None
+
+
+def test_maybe_recover_unrecognised_error_code():
+  llm_response = _make_rejected_response('SAFETY', 'print(42)')
+  assert _maybe_recover_from_api_rejection(llm_response) is False
+  assert llm_response.content is None
+
+
+def test_maybe_recover_no_error_code():
+  llm_response = LlmResponse(content=None, error_code=None, error_message=None)
+  assert _maybe_recover_from_api_rejection(llm_response) is False
+
+
+def test_maybe_recover_unparseable_message():
+  llm_response = LlmResponse(
+      content=None,
+      error_code='UNEXPECTED_TOOL_CALL',
+      error_message='some completely different message',
+  )
+  assert _maybe_recover_from_api_rejection(llm_response) is False