Add random rate limit responses to embedded webserver

b2listen/b2listen.py

@@ -193,7 +193,13 @@ def parse_args() -> argparse.Namespace:
     group.add_argument('--url', type=str,
                        help=f'Local webserver URL, for example: "http://localhost:8080")')  # noqa
     group.add_argument('--run-server', action='store_true',
-                       help=f'Run the embedded HTTP server')
+                       help=f'Run the embedded webserver')
+
+    run_server = parser_listen.add_argument_group(description='To simulate rate limiting in the embedded webserver:')
+    run_server.add_argument('--rate-limit-frequency', type=int, required=False,
+                              help='Respond with "429 Too Many Requests" for this percentage of notifications.')
+    run_server.add_argument('--retry-after', type=int, required=False,
+                            help='Value for the "Retry-After" header in a rate limit response.')
 
     use_existing = parser_listen.add_argument_group(description='To use an existing Event Notification rule:')
     use_existing.add_argument('--rule-name', type=str, required=False,
@@ -320,7 +326,8 @@ def run_cloudflared(command: str, loglevel: str, service_url: str, label: str, u
 
 def listen(args: argparse.Namespace):
     if args.run_server:
-        http_server = Server(interface='localhost', port=0, daemon=True)
+        http_server = Server(interface='localhost', port=0, daemon=True,
+                             rate_limit_frequency=args.rate_limit_frequency, retry_after=args.retry_after)
         http_server.start()
         service_url = f'http://{http_server.interface}:{http_server.port}'  # noqa
     else:

b2listen/server.py

@@ -9,11 +9,15 @@
 
 From https://gist.github.com/mdonkers/63e115cc0c79b4f6b8b3a6b797e485c7
 """
+import random
+from http import HTTPStatus
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from sys import argv
 import logging
 from threading import Thread
 
+RETRY_AFTER = 'Retry-After'
+
 logging.basicConfig()
 logger = logging.getLogger('b2listen.server')
 
@@ -22,15 +26,19 @@
 
 
 class S(BaseHTTPRequestHandler):
-    def _set_response(self):
-        self.send_response(200)
-        self.send_header('Content-type', 'text/html')
+    rate_limit_frequency = 0
+    retry_after = 0
+
+    def _set_response(self, status_code):
+        self.send_response(status_code)
+        if status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            self.send_header(RETRY_AFTER, str(self.retry_after))
         self.end_headers()
 
     # noinspection PyPep8Naming
     def do_GET(self):
         logger.info("GET request,\nPath: %s\nHeaders:\n%s\n", str(self.path), str(self.headers))
-        self._set_response()
+        self._set_response(HTTPStatus.OK)
         self.wfile.write("GET request for {}".format(self.path).encode('utf-8'))
 
     # noinspection PyPep8Naming
@@ -40,19 +48,23 @@ def do_POST(self):
         logger.info("POST request,\nPath: %s\nHeaders:\n%s\nBody:\n%s\n",
                     str(self.path), str(self.headers), post_data.decode('utf-8'))
 
-        self._set_response()
+        status_code = HTTPStatus.OK if random.random() > (self.rate_limit_frequency / 100) \
+            else HTTPStatus.TOO_MANY_REQUESTS
+        self._set_response(status_code)
         self.wfile.write("POST request for {}".format(self.path).encode('utf-8'))
 
 
 class Server(Thread):
     def __init__(self, server_class=HTTPServer, handler_class=S, interface=DEFAULT_INTERFACE, port=DEFAULT_PORT,
-                 daemon=False):
+                 daemon=False, rate_limit_frequency=0, retry_after=0):
         super().__init__(daemon=daemon)
         server_address = (interface, port)
         # noinspection PyTypeChecker
         self.httpd = server_class(server_address, handler_class)
         self.interface = self.httpd.server_address[0]
         self.port = self.httpd.server_address[1]
+        S.rate_limit_frequency = rate_limit_frequency
+        S.retry_after = retry_after
 
     def run(self):
         logger.info(f'Starting HTTP server on {self.interface}:{self.port}')