def test_duplicate_registration_of_same_class(self):
'''
Try to register the same class twice
'''
class ValidType(ControlMessage):
message_type = 1
type_registry.register_type_class(ValidType)
the_class = type_registry.get_type_class(1)
self.assertEqual(the_class, ValidType)
type_registry.register_type_class(ValidType)
the_class = type_registry.get_type_class(1)
self.assertEqual(the_class, ValidType)
def test_bad_duplicate_registration(self):
'''
Try to register two classes for the same type
'''
class ValidType1(ControlMessage):
message_type = 1
class ValidType2(ControlMessage):
message_type = 1
type_registry.register_type_class(ValidType1)
the_class = type_registry.get_type_class(1)
self.assertEqual(the_class, ValidType1)
with self.assertRaisesRegexp(ValueError, 'bound'):
type_registry.register_type_class(ValidType2)
the_class = type_registry.get_type_class(1)
self.assertEqual(the_class, ValidType1)
def test_register_type_class(self):
'''
Normal registration of a type class
'''
class ValidType(ControlMessage):
message_type = 1
type_registry.register_type_class(ValidType)
the_class = type_registry.get_type_class(1)
self.assertEqual(the_class, ValidType)
def from_bytes(cls, bitstream):
'''
Look at the type of the message, instantiate the correct class and
let it parse the message.
'''
from pylisp.packet.lisp.control import type_registry
# Convert to ConstBitStream (if not already provided)
if not isinstance(bitstream, ConstBitStream):
if isinstance(bitstream, Bits):
bitstream = ConstBitStream(auto=bitstream)
else:
bitstream = ConstBitStream(bytes=bitstream)
# Peek at the bitstream to see which type it is
type_nr = bitstream.peek('uint:4')
# Look for the right class
type_class = type_registry.get_type_class(type_nr)
if not type_class:
raise ValueError("Can't handle message type {0}".format(type_nr))
# Let the specific class handle it from now on
return type_class.from_bytes(bitstream)
