Add CLI arg to benchmark script to select a backend

1 files changed, 55 insertions, 14 deletions

diff --git a/examples/benchmark.py b/examples/benchmark.py
index b37ffef..6992b9a 100755
--- a/examples/benchmark.py
+++ b/examples/benchmark.py
@@ -1,46 +1,75 @@
 #!/usr/bin/env python
 """
 An example of benchmarking cache write speeds with semi-randomized response content
+
+Usage:
+```
+python benchmark.py -b <backend name>
+```
 """
+from argparse import ArgumentParser
 from os import urandom
 from random import random
 from time import perf_counter as time
 
 import requests
+from rich import print
+from rich.progress import Progress
 
 from requests_cache import CachedResponse, CachedSession
 
 BASE_RESPONSE = requests.get('https://httpbin.org/get')
-CACHE_NAME = 'rubbish_bin.sqlite'
+CACHE_NAME = 'rubbish_bin'
+WARMUP_ITERATIONS = 100
 ITERATIONS = 5000
-MAX_RESPONSE_SIZE = 1024 * 1024
+MAX_RESPONSE_SIZE = 1024 * 350
 
+# Defaults for DynamoDB
+AWS_OPTIONS = {
+    'endpoint_url': 'http://localhost:8000',
+    'region_name': 'us-east-1',
+    'aws_access_key_id': 'placeholder',
+    'aws_secret_access_key': 'placeholder',
+}
 
-def test_write_speed(session):
-    print(f'Testing write speed over {ITERATIONS} iterations')
-    start = time()
 
-    for i in range(ITERATIONS):
+def test_write_speed(session):
+    for i in range(WARMUP_ITERATIONS):
         new_response = get_randomized_response(i)
         session.cache.save_response(new_response)
 
+    with Progress() as progress:
+        task = progress.add_task('[cyan]Testing write speed...', total=ITERATIONS)
+        start = time()
+
+        for i in range(ITERATIONS):
+            new_response = get_randomized_response(i)
+            session.cache.save_response(new_response)
+            progress.update(task, advance=1)
+
     elapsed = time() - start
     avg = (elapsed / ITERATIONS) * 1000
-    print(f'Elapsed: {elapsed:.3f} (avg {avg:.3f}ms per write)')
+    print(f'[cyan]Elapsed: [green]{elapsed:.3f}[/] seconds (avg [green]{avg:.3f}[/] ms per write)')
 
 
 def test_read_speed(session):
-    print(f'Testing read speed over {ITERATIONS} iterations')
     keys = list(session.cache.responses.keys())
-    start = time()
-
-    for i in range(ITERATIONS):
-        key = keys[i % ITERATIONS]
+    for i in range(WARMUP_ITERATIONS):
+        key = keys[i % len(keys)]
         session.cache.get_response(key)
 
+    with Progress() as progress:
+        task = progress.add_task('[cyan]Testing read speed...', total=ITERATIONS)
+        start = time()
+
+        for i in range(ITERATIONS):
+            key = keys[i % len(keys)]
+            session.cache.get_response(key)
+            progress.update(task, advance=1)
+
     elapsed = time() - start
     avg = (elapsed / ITERATIONS) * 1000
-    print(f'Elapsed: {elapsed:.3f} (avg {avg:.3f}ms per read)')
+    print(f'[cyan]Elapsed: [green]{elapsed:.3f}[/] seconds (avg [green]{avg:.3f}[/] ms per read)')
 
 
 def get_randomized_response(i=0):
@@ -53,6 +82,18 @@ def get_randomized_response(i=0):
 
 
 if __name__ == '__main__':
-    session = CachedSession(CACHE_NAME, backend='sqlite')
+    parser = ArgumentParser()
+    parser.add_argument('-b', '--backend', default='sqlite')
+    args = parser.parse_args()
+    print(f'[cyan]Benchmarking {args.backend} backend')
+
+    kwargs = {}
+    if args.backend == 'dynamodb':
+        kwargs = AWS_OPTIONS
+    elif args.backend == 'sqlite-memory':
+        args.backend = 'sqlite'
+        kwargs = {'use_memory': True}
+
+    session = CachedSession(CACHE_NAME, backend=args.backend, **kwargs)
     test_write_speed(session)
     test_read_speed(session)