github
/
apprise
mirror of https://github.com/caronc/apprise
1
0
Fork 0
Code Issues Releases Wiki Activity

Discord Rate Limiting (429 Error code) handling (#901)

pull/902/head
Chris Caron 2023-07-09 14:38:40 -04:00 committed by GitHub
parent 78f5693382
commit a94700b0eb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 159 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
apprise/plugins/NotifyDiscord.py
View File

@ -50,6 +50,9 @@
import re
import requests
from json import dumps
from datetime import timedelta
from datetime import datetime
from datetime import timezone
from .NotifyBase import NotifyBase
from ..common import NotifyImageSize
@ -84,6 +87,17 @@ class NotifyDiscord(NotifyBase):
# Allows the user to specify the NotifyImageSize object
image_size = NotifyImageSize.XY_256
# Discord is kind enough to return how many more requests we're allowed to
# continue to make within it's header response as:
# X-RateLimit-Reset: The epoc time (in seconds) we can expect our
# rate-limit to be reset.
# X-RateLimit-Remaining: an integer identifying how many requests we're
# still allow to make.
request_rate_per_sec = 0
# Taken right from google.auth.helpers:
clock_skew = timedelta(seconds=10)
# The maximum allowable characters allowed in the body per message
body_maxlen = 2000
@ -215,6 +229,12 @@ class NotifyDiscord(NotifyBase):
# dynamically generated avatar url images
self.avatar_url = avatar_url
# For Tracking Purposes
self.ratelimit_reset = datetime.now(timezone.utc).replace(tzinfo=None)
# Default to 1.0
self.ratelimit_remaining = 1.0
return
def send(self, body, title='', notify_type=NotifyType.INFO, attach=None,
@ -343,7 +363,8 @@ class NotifyDiscord(NotifyBase):
# Otherwise return
return True
def _send(self, payload, attach=None, params=None, **kwargs):
def _send(self, payload, attach=None, params=None, rate_limit=1,
**kwargs):
"""
Wrapper to the requests (post) object
"""
@ -365,8 +386,25 @@ class NotifyDiscord(NotifyBase):
))
self.logger.debug('Discord Payload: %s' % str(payload))
# Always call throttle before any remote server i/o is made
self.throttle()
# By default set wait to None
wait = None
if self.ratelimit_remaining <= 0.0:
# Determine how long we should wait for or if we should wait at
# all. This isn't fool-proof because we can't be sure the client
# time (calling this script) is completely synced up with the
# Gitter server. One would hope we're on NTP and our clocks are
# the same allowing this to role smoothly:
now = datetime.now(timezone.utc).replace(tzinfo=None)
if now < self.ratelimit_reset:
# We need to throttle for the difference in seconds
wait = abs(
(self.ratelimit_reset - now + self.clock_skew)
.total_seconds())
# Always call throttle before any remote server i/o is made;
self.throttle(wait=wait)
# Perform some simple error checking
if isinstance(attach, AttachBase):
@ -401,6 +439,22 @@ class NotifyDiscord(NotifyBase):
verify=self.verify_certificate,
timeout=self.request_timeout,
)
# Handle rate limiting (if specified)
try:
# Store our rate limiting (if provided)
self.ratelimit_remaining = \
float(r.headers.get(
'X-RateLimit-Remaining'))
self.ratelimit_reset = datetime.fromtimestamp(
int(r.headers.get('X-RateLimit-Reset')),
timezone.utc).replace(tzinfo=None)
except (TypeError, ValueError):
# This is returned if we could not retrieve this
# information gracefully accept this state and move on
pass
if r.status_code not in (
requests.codes.ok, requests.codes.no_content):
@ -408,6 +462,20 @@ class NotifyDiscord(NotifyBase):
status_str = \
NotifyBase.http_response_code_lookup(r.status_code)
if r.status_code == requests.codes.too_many_requests \
and rate_limit > 0:
# handle rate limiting
self.logger.warning(
'Discord rate limiting in effect; '
'blocking for %.2f second(s)',
self.ratelimit_remaining)
# Try one more time before failing
return self._send(
payload=payload, attach=attach, params=params,
rate_limit=rate_limit - 1, **kwargs)
self.logger.warning(
'Failed to send {}to Discord notification: '
'{}{}error={}.'.format(

89
test/test_plugin_discord.py
View File

@ -32,7 +32,8 @@
import os
from unittest import mock
from datetime import datetime, timedelta
from datetime import timezone
import pytest
import requests
@ -182,6 +183,11 @@ def test_plugin_discord_general(mock_post):
"""
# Turn off clock skew for local testing
NotifyDiscord.clock_skew = timedelta(seconds=0)
# Epoch time:
epoch = datetime.fromtimestamp(0, timezone.utc)
# Initialize some generic (but valid) tokens
webhook_id = 'A' * 24
webhook_token = 'B' * 64
@ -189,6 +195,12 @@ def test_plugin_discord_general(mock_post):
# Prepare Mock
mock_post.return_value = requests.Request()
mock_post.return_value.status_code = requests.codes.ok
mock_post.return_value.content = ''
mock_post.return_value.headers = {
'X-RateLimit-Reset': (
datetime.now(timezone.utc) - epoch).total_seconds(),
'X-RateLimit-Remaining': 1,
}
# Invalid webhook id
with pytest.raises(TypeError):
@ -208,10 +220,85 @@ def test_plugin_discord_general(mock_post):
webhook_id=webhook_id,
webhook_token=webhook_token,
footer=True, thumbnail=False)
assert obj.ratelimit_remaining == 1
# Test that we get a string response
assert isinstance(obj.url(), str) is True
# This call includes an image with it's payload:
assert obj.notify(
body='body', title='title', notify_type=NotifyType.INFO) is True
# Force a case where there are no more remaining posts allowed
mock_post.return_value.headers = {
'X-RateLimit-Reset': (
datetime.now(timezone.utc) - epoch).total_seconds(),
'X-RateLimit-Remaining': 0,
}
# This call includes an image with it's payload:
assert obj.notify(
body='body', title='title', notify_type=NotifyType.INFO) is True
# behind the scenes, it should cause us to update our rate limit
assert obj.send(body="test") is True
assert obj.ratelimit_remaining == 0
# This should cause us to block
mock_post.return_value.headers = {
'X-RateLimit-Reset': (
datetime.now(timezone.utc) - epoch).total_seconds(),
'X-RateLimit-Remaining': 10,
}
assert obj.send(body="test") is True
assert obj.ratelimit_remaining == 10
# Reset our variable back to 1
mock_post.return_value.headers = {
'X-RateLimit-Reset': (
datetime.now(timezone.utc) - epoch).total_seconds(),
'X-RateLimit-Remaining': 1,
}
# Handle cases where our epoch time is wrong
del mock_post.return_value.headers['X-RateLimit-Reset']
assert obj.send(body="test") is True
# Return our object, but place it in the future forcing us to block
mock_post.return_value.headers = {
'X-RateLimit-Reset': (
datetime.now(timezone.utc) - epoch).total_seconds() + 1,
'X-RateLimit-Remaining': 0,
}
obj.ratelimit_remaining = 0
assert obj.send(body="test") is True
# Test 429 error response
mock_post.return_value.status_code = requests.codes.too_many_requests
# The below will attempt a second transmission and fail (because we didn't
# set up a second post request to pass) :)
assert obj.send(body="test") is False
# Return our object, but place it in the future forcing us to block
mock_post.return_value.status_code = requests.codes.ok
mock_post.return_value.headers = {
'X-RateLimit-Reset': (
datetime.now(timezone.utc) - epoch).total_seconds() - 1,
'X-RateLimit-Remaining': 0,
}
assert obj.send(body="test") is True
# Return our limits to always work
obj.ratelimit_remaining = 1
# Return our headers to normal
mock_post.return_value.headers = {
'X-RateLimit-Reset': (
datetime.now(timezone.utc) - epoch).total_seconds(),
'X-RateLimit-Remaining': 1,
}
# This call includes an image with it's payload:
assert obj.notify(
body='body', title='title', notify_type=NotifyType.INFO) is True