diff --git a/src/google/adk/tools/openapi_tool/common/common.py b/src/google/adk/tools/openapi_tool/common/common.py index 3b9b6b2497d..716b728716c 100644 --- a/src/google/adk/tools/openapi_tool/common/common.py +++ b/src/google/adk/tools/openapi_tool/common/common.py @@ -227,8 +227,13 @@ def generate_return_doc(responses: Dict[str, Response]) -> str: # Only consider 2xx responses for return type hinting. # Returns the 2xx response with the smallest status code number and with - # content defined. - sorted_responses = sorted(responses.items(), key=lambda item: int(item[0])) + # content defined. Skip non-numeric keys such as 'default' and range + # patterns ('2XX', '4XX'), which are valid OpenAPI response keys but cannot + # be parsed as an int (they are filtered out by the 2xx check below anyway). + sorted_responses = sorted( + (item for item in responses.items() if item[0].isdigit()), + key=lambda item: int(item[0]), + ) qualified_response = next( filter( lambda r: r[0].startswith('2') and r[1].content, diff --git a/tests/unittests/tools/openapi_tool/common/test_common.py b/tests/unittests/tools/openapi_tool/common/test_common.py index 1dd3195071f..20bca63c88b 100644 --- a/tests/unittests/tools/openapi_tool/common/test_common.py +++ b/tests/unittests/tools/openapi_tool/common/test_common.py @@ -407,6 +407,38 @@ def test_generate_return_doc_contentful_response(self): == expected_doc ) + def test_generate_return_doc_with_default_key_does_not_crash(self): + responses = { + '200': { + 'description': 'Successful response', + 'content': {'application/json': {'schema': {'type': 'string'}}}, + }, + 'default': {'description': 'Unexpected error'}, + } + expected_doc = 'Returns (str): Successful response' + assert ( + PydocHelper.generate_return_doc(dict_to_responses(responses)) + == expected_doc + ) + + def test_generate_return_doc_with_range_key_does_not_crash(self): + responses = { + '4XX': {'description': 'Client error'}, + '200': { + 'description': 'Successful response', + 'content': {'application/json': {'schema': {'type': 'string'}}}, + }, + } + expected_doc = 'Returns (str): Successful response' + assert ( + PydocHelper.generate_return_doc(dict_to_responses(responses)) + == expected_doc + ) + + def test_generate_return_doc_only_non_numeric_keys_returns_empty(self): + responses = {'default': {'description': 'Only a default response'}} + assert PydocHelper.generate_return_doc(dict_to_responses(responses)) == '' + if __name__ == '__main__': pytest.main([__file__])