diff --git a/locust/contrib/fasthttp.py b/locust/contrib/fasthttp.py index e44a2d99ec..6ba1852f2e 100644 --- a/locust/contrib/fasthttp.py +++ b/locust/contrib/fasthttp.py @@ -6,7 +6,7 @@ from urllib.parse import urlparse, urlunparse from ssl import SSLError import time -from typing import Optional +from typing import Optional, Tuple, Dict, Any from http.cookiejar import CookieJar @@ -15,7 +15,8 @@ from geventhttpclient._parser import HTTPParseError from geventhttpclient.client import HTTPClientPool from geventhttpclient.useragent import UserAgent, CompatRequest, CompatResponse, ConnectionError -from geventhttpclient.response import HTTPConnectionClosed +from geventhttpclient.response import HTTPConnectionClosed, HTTPSocketPoolResponse +from geventhttpclient.header import Headers from locust.user import User from locust.exception import LocustError, CatchResponseError, ResponseError @@ -70,7 +71,7 @@ def __init__( self, environment: Environment, base_url: str, - user: "FastHttpUser", + user: Optional[User], insecure=True, client_pool: Optional[HTTPClientPool] = None, **kwargs, @@ -123,14 +124,22 @@ def _send_request_safe_mode(self, method, url, **kwargs): if hasattr(e, "response"): r = e.response else: - r = ErrorResponse() + safe_kwargs = kwargs or {} + req = self.client._make_request( + url, + method=method, + headers=safe_kwargs.get("headers", None), + payload=safe_kwargs.get("payload", None), + params=safe_kwargs.get("params", None), + ) + r = ErrorResponse(url=url, request=req) r.error = e return r def request( self, method: str, - path: str, + url: str, name: str = None, data: str = None, catch_response: bool = False, @@ -147,7 +156,7 @@ def request( Returns :py:class:`locust.contrib.fasthttp.FastResponse` object. :param method: method for the new :class:`Request` object. - :param path: Path that will be concatenated with the base host URL that has been specified. + :param url: path that will be concatenated with the base host URL that has been specified. Can also be a full URL, in which case the full URL will be requested, and the base host is ignored. :param name: (optional) An argument that can be specified to use as label in Locust's @@ -170,7 +179,7 @@ def request( content will not be accounted for in the request time that is reported by Locust. """ # prepend url with hostname unless it's already an absolute URL - url = self._build_url(path) + built_url = self._build_url(url) start_time = time.time() # seconds since epoch @@ -198,15 +207,15 @@ def request( start_perf_counter = time.perf_counter() # send request, and catch any exceptions - response = self._send_request_safe_mode(method, url, payload=data, headers=headers, **kwargs) + response = self._send_request_safe_mode(method, built_url, payload=data, headers=headers, **kwargs) request_meta = { "request_type": method, - "name": name or path, + "name": name or url, "context": context, "response": response, "exception": None, "start_time": start_time, - "url": url, # this is a small deviation from HttpSession, which gets the final (possibly redirected) URL + "url": built_url, # this is a small deviation from HttpSession, which gets the final (possibly redirected) URL } if not allow_redirects: @@ -242,32 +251,32 @@ def request( self.environment.events.request.fire(**request_meta) return response - def delete(self, path, **kwargs): - return self.request("DELETE", path, **kwargs) + def delete(self, url, **kwargs): + return self.request("DELETE", url, **kwargs) - def get(self, path, **kwargs): + def get(self, url, **kwargs): """Sends a GET request""" - return self.request("GET", path, **kwargs) + return self.request("GET", url, **kwargs) - def head(self, path, **kwargs): + def head(self, url, **kwargs): """Sends a HEAD request""" - return self.request("HEAD", path, **kwargs) + return self.request("HEAD", url, **kwargs) - def options(self, path, **kwargs): + def options(self, url, **kwargs): """Sends a OPTIONS request""" - return self.request("OPTIONS", path, **kwargs) + return self.request("OPTIONS", url, **kwargs) - def patch(self, path, data=None, **kwargs): + def patch(self, url, data=None, **kwargs): """Sends a POST request""" - return self.request("PATCH", path, data=data, **kwargs) + return self.request("PATCH", url, data=data, **kwargs) - def post(self, path, data=None, **kwargs): + def post(self, url, data=None, **kwargs): """Sends a POST request""" - return self.request("POST", path, data=data, **kwargs) + return self.request("POST", url, data=data, **kwargs) - def put(self, path, data=None, **kwargs): + def put(self, url, data=None, **kwargs): """Sends a PUT request""" - return self.request("PUT", path, data=data, **kwargs) + return self.request("PUT", url, data=data, **kwargs) class FastHttpUser(User): @@ -338,15 +347,35 @@ def __init__(self, environment): """ +class FastRequest(CompatRequest): + payload: Optional[str] = None + + @property + def body(self) -> Optional[str]: + return self.payload + + class FastResponse(CompatResponse): - headers = None + headers: Optional[Headers] = None """Dict like object containing the response headers""" - _response = None + _response: Optional[HTTPSocketPoolResponse] = None encoding: Optional[str] = None """In some cases setting the encoding explicitly is needed. If so, do it before calling .text""" + request: Optional[FastRequest] = None + + def __init__( + self, + ghc_response: HTTPSocketPoolResponse, + request: Optional[FastRequest] = None, + sent_request: Optional[str] = None, + ): + super().__init__(ghc_response, request, sent_request) + + self.request = request + @property def text(self) -> Optional[str]: """ @@ -361,6 +390,16 @@ def text(self) -> Optional[str]: self.encoding = self.headers.get("content-type", "").partition("charset=")[2] or "utf-8" return str(self.content, self.encoding, errors="replace") + @property + def url(self) -> Optional[str]: + """ + Get "response" URL, which is the same as the request URL. This is a small deviation from HttpSession, which gets the final (possibly redirected) URL. + """ + if self.request is not None: + return self.request.url + + return None + def json(self) -> dict: """ Parses the response as json and returns a dict @@ -402,11 +441,16 @@ class ErrorResponse: that doesn't have a real Response object attached. E.g. a socket error or similar """ - headers = None + headers: Optional[Headers] = None content = None status_code = 0 - error = None - text = None + error: Optional[Exception] = None + text: Optional[str] = None + request: CompatRequest + + def __init__(self, url: str, request: CompatRequest): + self.url = url + self.request = request def raise_for_status(self): raise self.error @@ -414,6 +458,7 @@ def raise_for_status(self): class LocustUserAgent(UserAgent): response_type = FastResponse + request_type = FastRequest valid_response_codes = frozenset([200, 201, 202, 203, 204, 205, 206, 207, 208, 226, 301, 302, 303, 307]) def __init__(self, client_pool: Optional[HTTPClientPool] = None, **kwargs): diff --git a/locust/test/test_fasthttp.py b/locust/test/test_fasthttp.py index fed81a9f87..40f73211d1 100644 --- a/locust/test/test_fasthttp.py +++ b/locust/test/test_fasthttp.py @@ -25,12 +25,15 @@ def test_get(self): def test_connection_error(self): s = FastHttpSession(self.environment, "http://localhost:1", user=None) - r = s.get("/", timeout=0.1) + r = s.get("/", timeout=0.1, headers={"X-Test-Headers": "hello"}) self.assertEqual(r.status_code, 0) self.assertEqual(None, r.content) self.assertEqual(1, len(self.runner.stats.errors)) self.assertTrue(isinstance(r.error, ConnectionRefusedError)) self.assertTrue(isinstance(next(iter(self.runner.stats.errors.values())).error, ConnectionRefusedError)) + self.assertEqual(r.url, "http://localhost:1/") + self.assertEqual(r.request.url, r.url) + self.assertEqual(r.request.headers.get("X-Test-Headers", ""), "hello") def test_404(self): s = self.get_client() @@ -44,6 +47,8 @@ def test_204(self): self.assertEqual(204, r.status_code) self.assertEqual(1, self.runner.stats.get("/status/204", "GET").num_requests) self.assertEqual(0, self.runner.stats.get("/status/204", "GET").num_failures) + self.assertEqual(r.url, "http://127.0.0.1:%i/status/204" % self.port) + self.assertEqual(r.request.url, r.url) def test_streaming_response(self): """ @@ -88,6 +93,7 @@ def test_cookie(self): self.assertEqual(200, r.status_code) r = s.get("/get_cookie?name=testcookie") self.assertEqual("1337", r.content.decode()) + self.assertEqual("1337", r.text) def test_head(self): s = self.get_client() @@ -121,6 +127,8 @@ def test_json_payload(self): s = self.get_client() r = s.post("/request_method", json={"foo": "bar"}) self.assertEqual(200, r.status_code) + self.assertEqual(r.request.body, '{"foo": "bar"}') + self.assertEqual(r.request.headers.get("Content-Type", None), "application/json") def test_catch_response_fail_successful_request(self): s = self.get_client() @@ -312,7 +320,10 @@ class MyUser(FastHttpUser): host = "http://127.0.0.1:%i" % self.port locust = MyUser(self.environment) - self.assertEqual("hello", locust.client.get("/request_header_test", headers={"X-Header-Test": "hello"}).text) + r = locust.client.get("/request_header_test", headers={"X-Header-Test": "hello"}) + self.assertEqual("hello", r.text) + self.assertEqual("hello", r.headers.get("X-Header-Test", None)) + self.assertEqual("hello", r.request.headers.get("X-Header-Test", None)) def test_client_get(self): class MyUser(FastHttpUser): diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py index 0d0841e1bc..7d401d4204 100644 --- a/locust/test/test_runners.py +++ b/locust/test/test_runners.py @@ -1213,8 +1213,13 @@ def tick(self): 0, test_shape.get_current_user_count(), "Shape is not seeing stopped runner user count correctly" ) self.assertDictEqual(master.reported_user_classes_count, {"FixedUser1": 0, "FixedUser2": 0, "TestUser": 0}) - sleep(1.0) - self.assertEqual(STATE_STOPPED, master.state) + + try: + with gevent.Timeout(3.0): + while master.state != STATE_STOPPED: + sleep(0.1) + finally: + self.assertEqual(STATE_STOPPED, master.state) def test_distributed_shape_with_stop_timeout(self): """ diff --git a/locust/test/testcases.py b/locust/test/testcases.py index 1103245e9f..4c67e9b3c6 100644 --- a/locust/test/testcases.py +++ b/locust/test/testcases.py @@ -55,7 +55,11 @@ def request_method(): @app.route("/request_header_test") def request_header_test(): - return request.headers["X-Header-Test"] + x_header_test = request.headers["X-Header-Test"] + response = Response(x_header_test) + response.headers["X-Header-Test"] = x_header_test + + return response @app.route("/post", methods=["POST"])