def test_pack(self):
"""
测试打包大小
:return:
"""
packet = Packet(1, 'abcdef', 'abc', 'abc', 1, 1)
pack = packet.pack()
self.assertEqual(len(pack), PACKET_SIZE)
def test_transfer(self, mock_transfer, mock_commands):
"""
测试消息转发
:param mock_commands:
:return:
"""
packet = Packet(1, 'abcdef', 'abc', 'efc', 10, 0)
pack = packet.pack()
mock_commands.get.return_value = None
self.protocol.data_received(pack)
mock_commands.get.assert_called_with(10)
self.serverMock.loop.create_task.assert_called_with(mock_transfer())
def test_heartbeat(self, mock_commands):
"""
测试心跳包
:param mock_commands:
:return:
"""
packet = Packet(1, 'abcdef', 'abc', 'abc', Cmd.HEARTBEAT, 0)
pack = packet.pack()
self.protocol.data_received(pack)
# 测试Token一致
self.assertEqual(packet.token, self.serverMock.token)
mock_commands.get.assert_called_with(Cmd.HEARTBEAT)
def data_received(self, data):
self._buffer += data
while self._buffer:
if self.state == State.INIT:
if len(self._buffer) < PACKET_SIZE:
return
msg, self._buffer = self._buffer[:PACKET_SIZE], self._buffer[
PACKET_SIZE:]
try:
self.packet = Packet.unpack(msg)
except PacketException:
logging.debug(f'failed to parse message: {msg}')
self.state = State.INIT
self.transport.close()
return
if self.packet.token != self.server.token:
logging.debug(f"invalid token: {self.packet.token}")
self.transport.close()
logging.debug(
f"cmd: {self.packet.cmd} from {self.packet.from_} to {self.packet.to}"
)
self.state = State.REQUEST
if self.state == State.REQUEST:
if len(self._buffer) < self.packet.length:
return
msg, self._buffer = self._buffer[:self.packet.
length], self._buffer[
self.packet.length:]
self.packet.body = msg
logging.debug(f"content: {self.packet.body}")
command = commands.get(self.packet.cmd) or transfer
self.server.loop.create_task(command(self))
self.state = State.INIT
def test_unpack(self):
"""
测试解析包格式
:return:
"""
msg = b'\x00\x01abcdefabc\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00abc\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01'
packet = Packet.unpack(msg)
self.assertEqual(packet.ver, 1)
self.assertEqual(packet.from_, 'abc')
self.assertEqual(packet.to, 'abc')
self.assertEqual(packet.cmd, Cmd.HEARTBEAT)
self.assertEqual(packet.length, 1)
async def heartbeat(protocol):
"""
处理心跳包
:param protocol: Protocol
:return:
"""
from_ = protocol.packet.from_
try:
last_protocol = protocol.server.protocols[from_]
if last_protocol != protocol:
logging.debug(
f"close last connection with {from_} @ f{last_protocol.transport.get_extra_info('peername')}"
)
last_protocol.transport.close()
del protocol.server.protocols[from_]
except KeyError:
pass
protocol.server.protocols[from_] = protocol
protocol.heartbeat_at = time.time()
redis_key = f'{from_}_heartbeat'
redis_value = {
'uid': from_,
'ver': protocol.packet.ver,
'client_address': protocol.transport.get_extra_info('peername'),
'server_address': protocol.server.conn_ip,
}
await protocol.server.redis.execute('setex', redis_key,
protocol.server.expiration_time,
json.dumps(redis_value))
packet = Packet(protocol.packet.ver, protocol.server.token, from_, from_,
Cmd.REPLY, 0)
pack = packet.pack()
protocol.transport.write(pack)
def data_received(self, data):
self._buffer += data
if self.state == State.INIT:
if len(self._buffer) < PACKET_SIZE:
return
msg, self._buffer = self._buffer[:PACKET_SIZE], self._buffer[
PACKET_SIZE:]
self.packet = Packet.unpack(msg)
logging.info(
f"cmd: {self.packet.cmd} from {self.packet.from_} to {self.packet.to}"
)
self.state = State.REQUEST
if self.state == State.REQUEST:
if len(self._buffer) < self.packet.length:
return
msg, self._buffer = self._buffer[:self.packet.
length], self._buffer[self.packet.
length:]
self.packet.body = msg
logging.info(f"content: {self.packet.body}")
self.state = State.INIT
def data_received(self, data):
self.log.debug('received {!r}'.format(data))
packet = Packet.unpack(data)
print(packet.cmd)