|
1 | | -from typing import AsyncIterator |
| 1 | +from typing import AsyncIterator, Union |
2 | 2 | from unittest.mock import Mock |
3 | 3 |
|
4 | 4 | import pytest |
|
13 | 13 | _RunnableStreamAgent, # pyright: ignore [reportPrivateUsage] |
14 | 14 | _RunnableStreamOutputOnlyAgent, # pyright: ignore [reportPrivateUsage] |
15 | 15 | agent_wrapper, |
| 16 | + clean_docstring, |
16 | 17 | extract_fn_spec, |
17 | 18 | get_generic_args, |
18 | 19 | is_async_iterator, |
@@ -113,3 +114,62 @@ async def test_fn_stream_output_only(self, mock_api_client: Mock): |
113 | 114 | assert len(chunks) == 1 |
114 | 115 | assert isinstance(chunks[0], HelloTaskOutput) |
115 | 116 | assert chunks[0] == HelloTaskOutput(message="Hello, World!") |
| 117 | + |
| 118 | + |
| 119 | +@pytest.mark.parametrize( |
| 120 | + ("value", "expected"), |
| 121 | + [ |
| 122 | + # Empty docstrings |
| 123 | + ("", ""), |
| 124 | + (None, ""), |
| 125 | +
|
| 126 | + # Single line docstrings |
| 127 | + ("Hello world", "Hello world"), |
| 128 | + (" Hello world ", "Hello world"), |
| 129 | +
|
| 130 | + # Docstring with empty lines at start/end |
| 131 | + (""" |
| 132 | +
|
| 133 | + Hello world |
| 134 | +
|
| 135 | + """, "Hello world"), |
| 136 | +
|
| 137 | + # Multi-line docstring with indentation |
| 138 | + (""" |
| 139 | + First line |
| 140 | + Second line |
| 141 | + Indented line |
| 142 | + Last line |
| 143 | + """, "First line\nSecond line\n Indented line\nLast line"), |
| 144 | +
|
| 145 | + # Docstring with empty lines in between |
| 146 | + (""" |
| 147 | + First line |
| 148 | +
|
| 149 | + Second line |
| 150 | +
|
| 151 | + Third line |
| 152 | + """, "First line\n\nSecond line\n\nThird line"), |
| 153 | +
|
| 154 | + # Real-world example |
| 155 | + (""" |
| 156 | + Find the capital city of the country where the input city is located. |
| 157 | +
|
| 158 | + Guidelines: |
| 159 | + 1. First identify the country where the input city is located |
| 160 | + 2. Then provide the capital city of that country |
| 161 | + 3. Include an interesting historical or cultural fact about the capital |
| 162 | + 4. Be accurate and precise with geographical information |
| 163 | + 5. If the input city is itself the capital, still provide the information |
| 164 | + """, |
| 165 | + "Find the capital city of the country where the input city is located.\n\n" |
| 166 | + "Guidelines:\n" |
| 167 | + "1. First identify the country where the input city is located\n" |
| 168 | + "2. Then provide the capital city of that country\n" |
| 169 | + "3. Include an interesting historical or cultural fact about the capital\n" |
| 170 | + "4. Be accurate and precise with geographical information\n" |
| 171 | + "5. If the input city is itself the capital, still provide the information"), |
| 172 | + ], |
| 173 | +) |
| 174 | +def test_clean_docstring(value: Union[str, None], expected: str): |
| 175 | + assert clean_docstring(value) == expected |
0 commit comments