Merge "Fix serialization issue"

mistral/serialization.py

@@ -137,6 +137,15 @@ class PolymorphicSerializer(Serializer):
 
         self.serializers[key] = serializer
 
+    def unregister(self, entity_cls):
+        key = self._get_serialization_key(entity_cls)
+
+        if not key:
+            return
+
+        if key in self.serializers:
+            del self.serializers[key]
+
     def cleanup(self):
         self.serializers.clear()
 
@@ -154,6 +163,11 @@ class PolymorphicSerializer(Serializer):
 
         serializer = self.serializers.get(key)
 
+        if not serializer:
+            raise RuntimeError(
+                "Failed to find a serializer for the key: %s" % key
+            )
+
         result = {
             '__serial_key': key,
             '__serial_data': serializer.serialize(entity)
@@ -188,5 +202,9 @@ def register_serializer(entity_cls, serializer):
     get_polymorphic_serializer().register(entity_cls, serializer)
 
 
+def unregister_serializer(entity_cls):
+    get_polymorphic_serializer().unregister(entity_cls)
+
+
 def cleanup():
     get_polymorphic_serializer().cleanup()

mistral/tests/unit/test_serialization.py

@@ -42,7 +42,7 @@ class SerializationTest(base.BaseTest):
 
         serialization.register_serializer(MyClass, MyClassSerializer())
 
-        self.addCleanup(serialization.cleanup)
+        self.addCleanup(serialization.unregister_serializer, MyClass)
 
     def test_dict_based_serializer(self):
         obj = MyClass('a', 'b')