Skip to content

Commit

Permalink
fix: handle broken rate limit as best effort (#126)
Browse files Browse the repository at this point in the history
  • Loading branch information
ferhatelmas committed Aug 17, 2022
1 parent a63992f commit 683529d
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 3 deletions.
40 changes: 40 additions & 0 deletions stream_chat/tests/test_stream_response.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
from datetime import datetime, timezone

from stream_chat.types.stream_response import StreamResponse


class TestResponse:
def test_clean_header(self):
headers = {
"x-ratelimit-limit": "300",
"x-ratelimit-remaining": "299",
"x-ratelimit-reset": "1598806800",
}
response = StreamResponse({}, headers, 200)
assert response.rate_limit().limit == 300
assert response.rate_limit().remaining == 299
assert response.rate_limit().reset == datetime.fromtimestamp(
1598806800, timezone.utc
)

headers_2 = {
"x-ratelimit-limit": "300, 300",
"x-ratelimit-remaining": "299, 299",
"x-ratelimit-reset": "1598806800, 1598806800",
}
response = StreamResponse({}, headers_2, 200)
assert response.rate_limit().limit == 300
assert response.rate_limit().remaining == 299
assert response.rate_limit().reset == datetime.fromtimestamp(
1598806800, timezone.utc
)

headers_3 = {
"x-ratelimit-limit": "garbage",
"x-ratelimit-remaining": "garbage",
"x-ratelimit-reset": "garbage",
}
response = StreamResponse({}, headers_3, 200)
assert response.rate_limit().limit == 0
assert response.rate_limit().remaining == 0
assert response.rate_limit().reset == datetime.fromtimestamp(0, timezone.utc)
15 changes: 12 additions & 3 deletions stream_chat/types/stream_response.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,13 +38,22 @@ def __init__(
)
if limit and remaining and reset:
self.__rate_limit = RateLimitInfo(
limit=int(limit),
remaining=int(remaining),
reset=datetime.fromtimestamp(float(reset), timezone.utc),
limit=int(self._clean_header(limit)),
remaining=int(self._clean_header(remaining)),
reset=datetime.fromtimestamp(
float(self._clean_header(reset)), timezone.utc
),
)

super(StreamResponse, self).__init__(response_dict)

def _clean_header(self, header: str) -> int:
try:
values = (v.strip() for v in header.split(","))
return int(next(v for v in values if v))
except ValueError:
return 0

def rate_limit(self) -> Optional[RateLimitInfo]:
"""Returns the ratelimit info of your API operation."""
return self.__rate_limit
Expand Down

0 comments on commit 683529d

Please sign in to comment.