diff --git a/apprise/plugins/NotifyDiscord.py b/apprise/plugins/NotifyDiscord.py index fff76eef..408c86f5 100644 --- a/apprise/plugins/NotifyDiscord.py +++ b/apprise/plugins/NotifyDiscord.py @@ -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( diff --git a/test/test_plugin_discord.py b/test/test_plugin_discord.py index c6774d87..30c49977 100644 --- a/test/test_plugin_discord.py +++ b/test/test_plugin_discord.py @@ -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