Browse Source

Stresstests: All tests now work with the json serializer

Ask Solem 11 years ago
parent
commit
78ca518773
1 changed files with 51 additions and 2 deletions
  1. 51 2
      funtests/stress/stress/data.py

+ 51 - 2
funtests/stress/stress/data.py

@@ -1,9 +1,25 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import
 
+import json
+import sys
+
 from celery.utils.debug import humanbytes
+from celery.utils.imports import qualname
+
+type_registry = {}
+
+
+def json_reduce(obj, attrs):
+    return {'py/obj': {'type': qualname(obj), 'attrs': attrs}}
 
 
+def jsonable(cls):
+    type_registry[qualname(cls)] = cls.__from_json__
+    return cls
+
+
+@jsonable
 class Data(object):
 
     def __init__(self, label, data):
@@ -15,9 +31,42 @@ class Data(object):
             self.label, humanbytes(len(self.data)),
         )
 
+    def __to_json__(self):
+        return json_reduce(self, {'label': self.label, 'data': self.data})
+
+    @classmethod
+    def __from_json__(cls, label=None, data=None, **kwargs):
+        return cls(label, data)
+
     def __reduce__(self):
         return Data, (self.label, self.data)
     __unicode__ = __repr__ = __str__
 
-BIG = Data("BIG", 'x' * 2 ** 20 * 8)
-SMALL = Data("SMALL", 'e' * 1024)
+BIG = Data('BIG', 'x' * 2 ** 20 * 8)
+SMALL = Data('SMALL', 'e' * 1024)
+
+
+class JSONEncoder(json.JSONEncoder):
+
+    def default(self, obj):
+        try:
+            return super(JSONEncoder, self).default(obj)
+        except TypeError:
+            reducer = getattr(obj, '__to_json__', None)
+            if reducer:
+                return reducer()
+            raise
+
+
+def decode_hook(d):
+    try:
+        meta = d['py/obj']
+    except KeyError:
+        return d
+    type_registry[d['type']](**d['attrs'])
+
+
+def install_json():
+    json._default_encoder = JSONEncoder()
+    json._default_decoder.object_hook = decode_hook
+install_json()  # ugh, ugly but it's a test suite after all