def test_send_and_receive_message(self):
p = BasicProtocol()
p.transport = Transport()
session_factory = SessionFactory(ASession)
p.set_session_factory(session_factory)
self.assertFalse(p.send_message("123"))
msg = MessageHello()
self.assertFalse(p.send_message(msg))
p.connectionMade()
self.assertTrue(p.send_message(msg))
self.assertEqual(len(p.transport.buff), 1)
p.dataReceived(p.transport.buff[0])
self.assertIsInstance(p.session.msgs[0], MessageHello)
self.assertEquals(msg.timestamp, p.session.msgs[0].timestamp)
time.sleep(1)
msg = MessageHello()
self.assertNotEquals(msg.timestamp, p.session.msgs[0].timestamp)
self.assertTrue(p.send_message(msg))
self.assertEqual(len(p.transport.buff), 2)
db = DataBuffer()
db.append_string(p.transport.buff[1])
m = Message.deserialize(db)[0]
self.assertEqual(m.timestamp, msg.timestamp)
p.connectionLost()
self.assertNotIn('session', p.__dict__)
def test_get_resource(self):
conn = BasicProtocol()
conn.transport = Mock()
conn.server = Mock()
db = DataBuffer()
sess = TaskSession(conn)
sess.send = lambda m: db.append_string(m.serialize())
sess._can_send = lambda *_: True
sess.request_resource(str(uuid.uuid4()), TaskResourceHeader("tmp"))
assert Message.deserialize_message(db.buffered_data)
class BasicProtocol(SessionProtocol):
""" Connection-oriented basic protocol for twisted, support message serialization"""
def __init__(self):
self.opened = False
self.db = DataBuffer()
self.lock = Lock()
SessionProtocol.__init__(self)
def send_message(self, msg):
"""
Serialize and send message
:param Message msg: message to send
:return bool: return True if message has been send, False if an error has
"""
if not self.opened:
logger.error(msg)
logger.error("Send message failed - connection closed.")
return False
msg_to_send = self._prepare_msg_to_send(msg)
if msg_to_send is None:
return False
self.transport.getHandle()
self.transport.write(msg_to_send)
return True
def close(self):
"""
Close connection, after writing all pending (flush the write buffer and wait for producer to finish).
:return None:
"""
self.transport.loseConnection()
def close_now(self):
"""
Close connection ASAP, doesn't flush the write buffer or wait for the producer to finish
:return:
"""
self.opened = False
self.transport.abortConnection()
# Protocol functions
def connectionMade(self):
"""Called when new connection is successfully opened"""
SessionProtocol.connectionMade(self)
self.opened = True
def dataReceived(self, data):
"""Called when additional chunk of data is received from another peer"""
if not self._can_receive():
return None
if not self.session:
logger.warning("No session argument in connection state")
return None
self._interpret(data)
def connectionLost(self, reason=connectionDone):
"""Called when connection is lost (for whatever reason)"""
self.opened = False
if self.session:
self.session.dropped()
SessionProtocol.connectionLost(self, reason)
# Protected functions
def _prepare_msg_to_send(self, msg):
ser_msg = msg.serialize()
db = DataBuffer()
db.append_len_prefixed_string(ser_msg)
return db.read_all()
def _can_receive(self):
return self.opened and isinstance(self.db, DataBuffer)
def _interpret(self, data):
with self.lock:
self.db.append_string(data)
mess = self._data_to_messages()
# Interpret messages
if mess:
for m in mess:
self.session.interpret(m)
elif data:
logger.info("Deserialization of messages from {}:{} failed, maybe it's still "
"too short?".format(self.session.address, self.session.port))
def _data_to_messages(self):
return Message.deserialize(self.db)