diff options
| -rw-r--r-- | _pytest/test_http_check_ratelimited.py | 49 | ||||
| -rw-r--r-- | wee_slack.py | 10 |
2 files changed, 54 insertions, 5 deletions
diff --git a/_pytest/test_http_check_ratelimited.py b/_pytest/test_http_check_ratelimited.py new file mode 100644 index 0000000..948f247 --- /dev/null +++ b/_pytest/test_http_check_ratelimited.py @@ -0,0 +1,49 @@ +# -*- coding: utf-8 -*- + +from __future__ import print_function, unicode_literals +from textwrap import dedent + +from wee_slack import SlackRequest + + +def test_http_check_ratelimited_supports_multiple_headers(realish_eventrouter): + response = ( + dedent( + """ + HTTP/1.1 200 Connection established + + HTTP/2 200 + content-type: application/json; charset=utf-8 + + {"ok": true} + """ + ) + .strip() + .replace("\n", "\r\n") + ) + + request_metadata = SlackRequest(None, "", token="xoxp-1") + body, error = realish_eventrouter.http_check_ratelimited(request_metadata, response) + assert body == '{"ok": true}' + assert error == "" + + +def test_http_check_ratelimited_return_error_when_ratelimited(realish_eventrouter): + response = ( + dedent( + """ + HTTP/2 429 + content-type: application/json; charset=utf-8 + retry-after: 10 + + {"ok": false, "error": "ratelimited"} + """ + ) + .strip() + .replace("\n", "\r\n") + ) + + request_metadata = SlackRequest(None, "", token="xoxp-1") + body, error = realish_eventrouter.http_check_ratelimited(request_metadata, response) + assert body == "" + assert error == "ratelimited" diff --git a/wee_slack.py b/wee_slack.py index 84950a7..ed55ef0 100644 --- a/wee_slack.py +++ b/wee_slack.py @@ -668,19 +668,19 @@ class EventRouter(object): self.receive(message_json) def http_check_ratelimited(self, request_metadata, response): - headers_end_index = response.index("\r\n\r\n") - headers = response[:headers_end_index].split("\r\n") - http_status = headers[0].split(" ")[1] + parts = response.split("\r\n\r\nHTTP/") + last_header_part, body = parts[-1].split("\r\n\r\n", 1) + header_lines = last_header_part.split("\r\n") + http_status = header_lines[0].split(" ")[1] if http_status == "429": - for header in headers[1:]: + for header in header_lines[1:]: name, value = header.split(":", 1) if name.lower() == "retry-after": retry_after = int(value.strip()) request_metadata.retry_time = time.time() + retry_after return "", "ratelimited" - body = response[headers_end_index + 4 :] return body, "" def retry_request(self, request_metadata, data, return_code, err): |