Suggestion to Add "iter_lines" Support for "FastHttpUser" in Locust

[MasterKey-Pro]

Prerequisites

• I am using the latest version of Locust
• I am suggesting a new feature, not asking a question

Description

In testing AI services using HTTP Server-Sent Events (SSE), HttpUser allows handling of streaming data efficiently. However, using FastHttpUser for performance benefits isn't feasible due to its lack of streaming support.

Proposed Enhancement:

Implement iter_lines support in FastHttpUser to process streaming responses efficiently.
Maintain the performance benefits that FastHttpUser offers.
Benefits:

Expands FastHttpUser to handle real-time data and streaming APIs.
Simplifies testing by allowing users to utilize performance benefits without additional handling logic.
Provides consistent capabilities across user classes in Locust.
I believe this enhancement will increase Locust's versatility.

Thank you for considering this suggestion

[cyberw]

Hi @MasterKey-Pro !

I would love to have this feature added. Have you looked at implementing it? I guess the stuff already in gevenhttpclient (https://github.com/geventhttpclient/geventhttpclient/tree/2e9c0a0f73b18ddb83fe8e62c57123daee9112f5?tab=readme-ov-file#streaming) is not enough for your use case?

[MasterKey-Pro]

Hello @cyberw ！

Thank you so much for your attention and support regarding my feature suggestion. I truly appreciate the excellent work you have done on the project.

I noticed that geventhttpclient already provides the necessary capabilities to handle streaming responses effectively. However, by implementing the iter_lines functionality directly in FastHttpUser, we can offer a more integrated and user-friendly solution that aligns seamlessly with the existing user interface. This approach would enhance usability by allowing users to easily handle streaming data without additional setup or complex configurations.

That said, I am concerned that my modifications might affect other functionalities or introduce unforeseen issues within the project. To address this, I have attached my code suggestions below for further discussion and improvement:

from geventhttpclient import HTTPClient

# Original FastHttpSession class
class FastHttpSession(HTTPClient):
    pass

# Mixin to add line iteration functionality
class LineIteratorMixin:
    def iter_lines(self, url, **kwargs):
        response = self.get(url, **kwargs)
        response.raise_for_status()

        buffer = ''
        for chunk in response.iter_content(chunk_size=1024):
            buffer += chunk.decode('utf-8')
            while '\n' in buffer:
                line, buffer = buffer.split('\n', 1)
                yield line

# Extended session class that includes the mixin
class ExtendedFastHttpSession(FastHttpSession, LineIteratorMixin):
    pass

class FastHttpUser(User):
    def __init__(self, environment) -> None:
        super().__init__(environment)
        if self.host is None:
            raise LocustError(
                "You must specify the base host. Either in the host attribute in the User class, or on the command line using the --host option."
            )

        # Use the extended session class with the new feature
        self.client = ExtendedFastHttpSession(
            self.environment,
            base_url=self.host,
            network_timeout=self.network_timeout,
            connection_timeout=self.connection_timeout,
            max_redirects=self.max_redirects,
            max_retries=self.max_retries,
            insecure=self.insecure,
            concurrency=self.concurrency,
            user=self,
            client_pool=self.client_pool,
            ssl_context_factory=self.ssl_context_factory,
            headers=self.default_headers,
            proxy_host=self.proxy_host,
            proxy_port=self.proxy_port,
        )

# Example user class
class VibeMateUser(FastHttpUser):
    def __init__(self, environment) -> None:
        super().__init__(environment)

    @task
    def stream_data(self):
        url = "/your/streaming/endpoint"
        for line in self.client.iter_lines(url):
            print(f"Received line: {line}")

    @task
    def get_data(self):
        url = "/your/normal/endpoint"
        response = self.client.get(url)
        print(f"Received response: {response.text}")

I greatly appreciate your support and assistance, and I hope my suggestions can contribute positively to the project’s development. If there are any questions or if further discussion is needed, please feel free to contact me.

[cyberw]

Looks good! I wouldn't bother making it a mixin though, just add the iter_lines method to FastHttpSession.

Also:

• to make a PR, follow the instructions here: https://docs.locust.io/en/stable/developing-locust.html
• add iter_lines here so that it becomes visible in the docs:

  locust/docs/api.rst

  Line 22 in 944b7f8

  :members: wait_time, tasks, client, abstract, rest

• add a test case for it