diff --git a/locust/contrib/fasthttp.py b/locust/contrib/fasthttp.py index 9c0f7a9a31..aee368cc61 100644 --- a/locust/contrib/fasthttp.py +++ b/locust/contrib/fasthttp.py @@ -298,6 +298,17 @@ def get(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager """Sends a GET request""" return self.request("GET", url, **kwargs) + def iter_lines(self, url: str, **kwargs) -> Generator[str, None, None]: + """Sends a iter_lines request""" + response = self.request("GET", url, stream=True, **kwargs) + response.raise_for_status() + buffer = "" + for chunk in response.iter_content(chunk_size=1024, decode_content=True): + buffer += chunk.decode("utf-8") + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + yield line + def head(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse: """Sends a HEAD request""" return self.request("HEAD", url, **kwargs) diff --git a/locust/test/test_fasthttp.py b/locust/test/test_fasthttp.py index b14d293829..a22a3aea81 100644 --- a/locust/test/test_fasthttp.py +++ b/locust/test/test_fasthttp.py @@ -8,6 +8,7 @@ import socket import time from tempfile import NamedTemporaryFile +from unittest.mock import MagicMock import gevent from geventhttpclient.client import HTTPClientPool @@ -107,6 +108,38 @@ def test_streaming_response_catch_response(self): self.assertGreaterEqual(stats.avg_response_time, 0) self.assertLess(stats.avg_response_time, 250) + def test_iter_lines(self): + session = self.get_client() + + url = "/streaming/10" + + response_mock = MagicMock() + response_mock.iter_content = MagicMock( + return_value=iter( + [ + b"0\n", + b"1\n", + b"2\n", + b"3\n", + b"4\n", + b"5\n", + b"6\n", + b"7\n", + b"8\n", + b"9\n", + ] + ) + ) + response_mock.raise_for_status = MagicMock() + + session.request = MagicMock(return_value=response_mock) + + lines = list(session.iter_lines(url)) + + expected_lines = [f"{i}" for i in range(10)] + + self.assertEqual(lines, expected_lines) + def test_slow_redirect(self): s = self.get_client() url = "/redirect?url=/redirect&delay=0.5"