Merge pull request #2416 from locustio/make-FastHttpUser-use-requests…

…-encoding-detection

Make FastHttpUser use requests encoding detection

locust/contrib/fasthttp.py

@@ -22,6 +22,9 @@
 from geventhttpclient.response import HTTPConnectionClosed, HTTPSocketPoolResponse
 from geventhttpclient.header import Headers
 
+# borrow requests's content-type header parsing
+from requests.utils import get_encoding_from_headers
+
 from locust.user import User
 from locust.exception import LocustError, CatchResponseError, ResponseError
 from locust.env import Environment
@@ -460,7 +463,7 @@ def text(self) -> Optional[str]:
             if self.headers is None:
                 self.encoding = "utf-8"
             else:
-                self.encoding = self.headers.get("content-type", "").partition("charset=")[2] or "utf-8"
+                self.encoding = get_encoding_from_headers(self.headers)
         return str(self.content, self.encoding, errors="replace")
 
     @property

locust/test/test_fasthttp.py

@@ -423,6 +423,16 @@ class MyUser(FastHttpUser):
         locust = MyUser(self.environment)
         self.assertEqual(200, locust.client.head("/request_method").status_code)
 
+    def test_complex_content_type(self):
+        class MyUser(FastHttpUser):
+            host = "http://127.0.0.1:%i" % self.port
+
+        locust = MyUser(self.environment)
+
+        self.assertEqual("stuff", locust.client.get("/content_type_missing_charset").text)
+        self.assertEqual("stuff", locust.client.get("/content_type_regular").text)
+        self.assertEqual("stuff", locust.client.get("/content_type_with_extra_stuff").text)
+
     def test_log_request_name_argument(self):
         self.response = ""
 

locust/test/testcases.py

@@ -149,6 +149,27 @@ def rest():
     return request.json
 
 
+@app.route("/content_type_missing_charset")
+def content_type_missing_charset():
+    resp = make_response("stuff")
+    resp.headers["Content-Type"] = "Content-Type: application/json;"
+    return resp
+
+
+@app.route("/content_type_regular")
+def content_type_regular():
+    resp = make_response("stuff")
+    resp.headers["Content-Type"] = "Content-Type: application/json; charset=utf-8;"
+    return resp
+
+
+@app.route("/content_type_with_extra_stuff")
+def content_type_with_extra_stuff():
+    resp = make_response("stuff")
+    resp.headers["Content-Type"] = "Content-Type: application/json; charset=utf-8; api-version=3.0"
+    return resp
+
+
 class LocustTestCase(unittest.TestCase):
     """
     Test case class that restores locust.events.EventHook listeners on tearDown, so that it is