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)
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)