Merge pull request #2083 from mgor/feature/fasthttp_improvements

FastHttpUser improvements (including a rename of parameter "url" to "path")

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,18 +441,24 @@ 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
 
 
 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):

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):

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):
         """

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"])