Fix generate_test_db script, and add a basic script for testing cache write speeds

3 files changed, 83 insertions, 40 deletions

diff --git a/examples/benchmark.py b/examples/benchmark.py
new file mode 100755
index 0000000..1d0fd8e
--- /dev/null
+++ b/examples/benchmark.py
@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+"""An example of benchmarking cache write speeds with semi-randomized response content"""
+from os import urandom
+from random import random
+from time import perf_counter as time
+
+import requests
+
+from requests_cache import CachedResponse, CachedSession
+
+BASE_RESPONSE = requests.get('https://httpbin.org/get')
+CACHE_NAME = 'rubbish_bin.sqlite'
+ITERATIONS = 5000
+MAX_RESPONSE_SIZE = 1024 * 6
+
+
+def test_write_speed(session):
+    start = time()
+    for i in range(ITERATIONS):
+        new_response = get_randomized_response(i)
+        session.cache.save_response(new_response)
+
+    elapsed = time() - start
+    avg = (elapsed / ITERATIONS) * 1000
+    print(f'Elapsed: {elapsed:.3f} (avg {avg:.3f}ms per write)')
+
+
+def get_randomized_response(i=0):
+    """Get a response with randomized content"""
+    new_response = CachedResponse.from_response(BASE_RESPONSE)
+    n_bytes = int(random() * MAX_RESPONSE_SIZE)
+    new_response._content = urandom(n_bytes)
+    new_response.request.url += f'/response_{i}'
+    return new_response
+
+
+if __name__ == '__main__':
+    session = CachedSession(CACHE_NAME, backend='sqlite')
+    test_write_speed(session)
diff --git a/examples/generate_test_db.py b/examples/generate_test_db.py
index e3ce60f..6137fff 100644..100755
--- a/examples/generate_test_db.py
+++ b/examples/generate_test_db.py
@@ -3,18 +3,21 @@
 This is useful, for example, for reproducing issues that only occur with large caches.
 """
 import logging
+from datetime import datetime, timedelta
 from os import urandom
 from os.path import getsize
 from random import random
 from time import perf_counter as time
 
+import requests
 from rich.progress import Progress
 
 from requests_cache import ALL_METHODS, CachedSession, CachedResponse
-from requests_cache.response import format_file_size
-from tests.conftest import HTTPBIN_FORMATS, HTTPBIN_METHODS, httpbin
+from requests_cache.models.response import format_file_size
+from tests.conftest import HTTPBIN_FORMATS, HTTPBIN_METHODS
 
 
+BASE_RESPONSE = requests.get('https://httpbin.org/get')
 CACHE_NAME = 'rubbish_bin.sqlite'
 HTTPBIN_EXTRA_ENDPOINTS = [
     'anything',
@@ -41,37 +44,45 @@ class InvalidResponse(CachedResponse):
 
 def populate_cache(progress, task):
     session = CachedSession(CACHE_NAME, backend='sqlite', allowable_methods=ALL_METHODS)
+    n_previous_responses = len(session.cache.responses)
 
     # Cache a variety of different response formats, which may result in different behavior
-    urls = [('GET', httpbin(endpoint)) for endpoint in HTTPBIN_FORMATS + HTTPBIN_EXTRA_ENDPOINTS]
-    urls += [(method, httpbin(method.lower())) for method in HTTPBIN_METHODS]
+    urls = [
+        ('GET', 'https://httpbin.org/{endpoint}')
+        for endpoint in HTTPBIN_FORMATS + HTTPBIN_EXTRA_ENDPOINTS
+    ]
+    urls += [(method, 'https://httpbin.org/{method.lower()}') for method in HTTPBIN_METHODS]
     for method, url in urls:
         session.request(method, url)
         progress.update(task, advance=1)
 
     # Cache a large number of responses with randomized response content, which will expire at random times
-    response = session.get(httpbin('get'))
     with session.cache.responses.bulk_commit():
         for i in range(N_RESPONSES):
-            new_response = CachedResponse(response)
-            n_bytes = int(random() * MAX_RESPONSE_SIZE)
-            new_response._content = urandom(n_bytes)
-
-            new_response.request.url += f'/response_{i}'
-            expire_after = random() * MAX_EXPIRE_AFTER
-            session.cache.save_response(new_response, expire_after=expire_after)
+            new_response = get_randomized_response(i + n_previous_responses)
+            expires = datetime.now() + timedelta(seconds=random() * MAX_EXPIRE_AFTER)
+            session.cache.save_response(new_response, expires=expires)
             progress.update(task, advance=1)
 
     # Add some invalid responses
     with session.cache.responses.bulk_commit():
         for i in range(N_INVALID_RESPONSES):
-            new_response = InvalidResponse(response)
+            new_response = InvalidResponse.from_response(BASE_RESPONSE)
             new_response.request.url += f'/invalid_response_{i}'
             key = session.cache.create_key(new_response.request)
             session.cache.responses[key] = new_response
             progress.update(task, advance=1)
 
 
+def get_randomized_response(i=0):
+    """Get a response with randomized content"""
+    new_response = CachedResponse.from_response(BASE_RESPONSE)
+    n_bytes = int(random() * MAX_RESPONSE_SIZE)
+    new_response._content = urandom(n_bytes)
+    new_response.request.url += f'/response_{i}'
+    return new_response
+
+
 def remove_expired_responses(expire_after=None):
     logger.setLevel('DEBUG')
     session = CachedSession(CACHE_NAME)
diff --git a/tests/benchmark_serializers.py b/tests/benchmark_serializers.py
index b587a76..4b786fd 100644
--- a/tests/benchmark_serializers.py
+++ b/tests/benchmark_serializers.py
@@ -42,7 +42,7 @@ except ImportError:
 sys.path.insert(0, os.path.abspath('..'))
 
 from requests_cache import CachedSession
-from requests_cache.serializers import BaseSerializer, BSONSerializer, JSONSerializer, PickleSerializer
+from requests_cache.serializers import CattrStage, bson_serializer, json_serializer, pickle_serializer
 
 ITERATIONS = 10000
 
@@ -56,50 +56,43 @@ r = session.get('https://httpbin.org/get?x=y')
 
 
 def run_pickle():
-    run_serialize_deserialize('pickle', pickle.dumps, pickle.loads)
+    run_serialize_deserialize('pickle', pickle)
 
 
 def run_cattrs():
-    s = PickleSerializer()
-    run_serialize_deserialize('cattrs', s.unstructure, s.structure)
+    run_serialize_deserialize('cattrs', CattrStage)
 
 
 def run_cattrs_pickle():
-    s = PickleSerializer()
-    run_serialize_deserialize('cattrs+pickle', s.dumps, s.loads)
+    run_serialize_deserialize('cattrs+pickle', pickle_serializer)
 
 
-def run_cattrs_json():
-    s = BaseSerializer(converter_factory=make_converter)
-    run_serialize_deserialize(
-        'cattrs+json',
-        lambda obj: json.dumps(s.unstructure(obj)),
-        lambda obj: s.structure(json.loads(obj)),
-    )
+# def run_cattrs_json():
+#     s = CattrStage(converter_factory=make_converter)
+#     run_serialize_deserialize(
+#         'cattrs+json',
+#         lambda obj: json.dumps(s.unstructure(obj)),
+#         lambda obj: s.structure(json.loads(obj)),
+#     )
 
 
 def run_cattrs_ujson():
-    s = BaseSerializer(converter_factory=make_converter)
-    run_serialize_deserialize(
-        'cattrs+ujson',
-        lambda obj: ujson.dumps(s.unstructure(obj)),
-        lambda obj: s.structure(ujson.loads(obj)),
-    )
+    s = CattrStage(converter_factory=make_converter)
+    run_serialize_deserialize('cattrs+ujson', json_serializer)
 
 
 def run_cattrs_bson():
-    s = BSONSerializer()
-    run_serialize_deserialize('cattrs+bson', s.dumps, s.loads)
+    run_serialize_deserialize('cattrs+bson', bson_serializer)
 
 
-def run_serialize_deserialize(module, serialize, deserialize):
+def run_serialize_deserialize(module, serializer):
     start = time()
-    serialized = [serialize(r) for i in range(ITERATIONS)]
-    print(f'{module}.{serialize.__name__} x{ITERATIONS}: {time() - start:.3f}')
+    serialized = [serializer.dumps(r) for i in range(ITERATIONS)]
+    print(f'{module}.{serializer.__name__}.loads() x{ITERATIONS}: {time() - start:.3f}')
 
     start = time()
-    deserialized = [deserialize(obj) for obj in serialized]
-    print(f'{module}.{deserialize.__name__} x{ITERATIONS}: {time() - start:.3f}')
+    deserialized = [serializer.loads(obj) for obj in serialized]
+    print(f'{module}.{serializer.__name__}.dumps() x{ITERATIONS}: {time() - start:.3f}')
 
 
 if __name__ == '__main__':
@@ -108,7 +101,7 @@ if __name__ == '__main__':
     run_pickle()
     run_cattrs()
     run_cattrs_pickle()
-    run_cattrs_json()
+    # run_cattrs_json()
     run_cattrs_ujson()
     run_cattrs_bson()