diff options
Diffstat (limited to 'deluge/tests/test_transfer.py')
| -rw-r--r-- | deluge/tests/test_transfer.py | 398 |
1 files changed, 398 insertions, 0 deletions
diff --git a/deluge/tests/test_transfer.py b/deluge/tests/test_transfer.py new file mode 100644 index 0000000..92e349b --- /dev/null +++ b/deluge/tests/test_transfer.py @@ -0,0 +1,398 @@ +# +# Copyright (C) 2012 Bro <bro.development@gmail.com> +# +# This file is part of Deluge and is licensed under GNU General Public License 3.0, or later, with +# the additional special exception to link portions of this program with the OpenSSL library. +# See LICENSE for more details. +# + +import base64 + +import pytest +import rencode + +import deluge.log +from deluge.transfer import DelugeTransferProtocol + +deluge.log.setup_logger('none') + + +class TransferTestClass(DelugeTransferProtocol): + def __init__(self): + DelugeTransferProtocol.__init__(self) + self.transport = self + self.messages_out = [] + self.messages_in = [] + self.packet_count = 0 + + def write(self, message): + """ + Called by DelugeTransferProtocol class + This simulates the write method of the self.transport in DelugeTransferProtocol. + """ + self.messages_out.append(message) + + def message_received(self, message): + """ + This method overrides message_received is DelugeTransferProtocol and is + called with the complete message as it was sent by DelugeRPCProtocol + """ + self.messages_in.append(message) + + def get_messages_out_joined(self): + return b''.join(self.messages_out) + + def get_messages_in(self): + return self.messages_in + + def data_received_old_protocol(self, data): + """ + This is the original method logic (as close as possible) for handling data receival on the client + + :param data: a zlib compressed string encoded with rencode. + + """ + import zlib + + print('\n=== New Data Received ===\nBytes received:', len(data)) + + if self._buffer: + # We have some data from the last dataReceived() so lets prepend it + print('Current buffer:', len(self._buffer) if self._buffer else '0') + data = self._buffer + data + self._buffer = None + + self.packet_count += 1 + self._bytes_received += len(data) + + while data: + print('\n-- Handle packet data --') + + print('Bytes received:', self._bytes_received) + print('Current data:', len(data)) + + if self._message_length == 0: + # handle_new_message uses _buffer so set data to _buffer. + self._buffer = data + self._handle_new_message() + data = self._buffer + self._buffer = None + self.packet_count = 1 + print('New message of length:', self._message_length) + + dobj = zlib.decompressobj() + try: + request = rencode.loads(dobj.decompress(data)) + print('Successfully loaded message', end=' ') + print( + ' - Buffer length: %d, data length: %d, unused length: %d' + % ( + len(data), + len(data) - len(dobj.unused_data), + len(dobj.unused_data), + ) + ) + print('Packet count:', self.packet_count) + except Exception as ex: + # log.debug('Received possible invalid message (%r): %s', data, e) + # This could be cut-off data, so we'll save this in the buffer + # and try to prepend it on the next dataReceived() + self._buffer = data + print( + 'Failed to load buffer (size %d): %s' % (len(self._buffer), str(ex)) + ) + return + else: + data = dobj.unused_data + self._message_length = 0 + + self.message_received(request) + + +class TestDelugeTransferProtocol: + @pytest.fixture(autouse=True) + def set_up(self): + """ + The expected messages corresponds to the test messages (msg1, msg2) after they've been processed + by DelugeTransferProtocol.send, which means that they've first been encoded with rencode, + and then compressed with zlib. + The expected messages are encoded in base64 to easily including it here in the source. + So before comparing the results with the expected messages, the expected messages must be decoded, + or the result message be encoded in base64. + + """ + self.transfer = TransferTestClass() + self.msg1 = ( + 0, + 1, + {'key_int': 1242429423}, + {'key_str': b'some string'}, + {'key_bool': True}, + ) + self.msg2 = ( + 2, + 3, + {'key_float': 12424.29423}, + {'key_unicode': 'some string'}, + {'key_dict_with_tuple': {'key_tuple': (1, 2, 3)}}, + {'keylist': [4, '5', 6.7]}, + ) + + self.msg1_expected_compressed_base64 = ( + b'AQAAADF4nDvKwJjenp1aGZ+ZV+Lgxfv9PYRXXFLU' + b'XZyfm6oAZGTmpad3gAST8vNznAEAJhSQ' + ) + self.msg2_expected_compressed_base64 = ( + b'AQAAAF14nDvGxJzemZ1aGZ+Wk59Y4uTmpKib3g3il+ZlJuenpH' + b'YX5+emKhSXFGXmpadPBkmkZCaXxJdnlmTEl5QW5KRCdIOZhxmB' + b'hrUDuTmZxSWHWRpNnRyupaUBAHYlJxI=' + ) + + def test_send_one_message(self): + """ + Send one message and test that it has been sent correctoly to the + method 'write' in self.transport. + + """ + self.transfer.transfer_message(self.msg1) + # Get the data as sent by DelugeTransferProtocol + messages = self.transfer.get_messages_out_joined() + base64_encoded = base64.b64encode(messages) + assert base64_encoded == self.msg1_expected_compressed_base64 + + def test_receive_one_message(self): + """ + Receive one message and test that it has been sent to the + method 'message_received'. + + """ + self.transfer.dataReceived( + base64.b64decode(self.msg1_expected_compressed_base64) + ) + # Get the data as sent by DelugeTransferProtocol + messages = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg1) == rencode.dumps(messages) + + def test_receive_old_message(self): + """ + Receive an old message (with no header) and verify that the data is discarded. + + """ + self.transfer.dataReceived(rencode.dumps(self.msg1)) + assert len(self.transfer.get_messages_in()) == 0 + assert self.transfer._message_length == 0 + assert len(self.transfer._buffer) == 0 + + def test_receive_two_concatenated_messages(self): + """ + This test simply concatenates two messsages (as they're sent over the network), + and lets DelugeTransferProtocol receive the data as one string. + + """ + two_concatenated = base64.b64decode( + self.msg1_expected_compressed_base64 + ) + base64.b64decode(self.msg2_expected_compressed_base64) + self.transfer.dataReceived(two_concatenated) + + # Get the data as sent by DelugeTransferProtocol + message1 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg1) == rencode.dumps(message1) + message2 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg2) == rencode.dumps(message2) + + def test_receive_three_messages_in_parts(self): + """ + This test concatenates three messsages (as they're sent over the network), + and lets DelugeTransferProtocol receive the data in multiple parts. + + """ + msg_bytes = ( + base64.b64decode(self.msg1_expected_compressed_base64) + + base64.b64decode(self.msg2_expected_compressed_base64) + + base64.b64decode(self.msg1_expected_compressed_base64) + ) + packet_size = 40 + + one_message_byte_count = len( + base64.b64decode(self.msg1_expected_compressed_base64) + ) + two_messages_byte_count = one_message_byte_count + len( + base64.b64decode(self.msg2_expected_compressed_base64) + ) + three_messages_byte_count = two_messages_byte_count + len( + base64.b64decode(self.msg1_expected_compressed_base64) + ) + + for d in self.receive_parts_helper(msg_bytes, packet_size): + bytes_received = self.transfer.get_bytes_recv() + + if bytes_received >= three_messages_byte_count: + expected_msgs_received_count = 3 + elif bytes_received >= two_messages_byte_count: + expected_msgs_received_count = 2 + elif bytes_received >= one_message_byte_count: + expected_msgs_received_count = 1 + else: + expected_msgs_received_count = 0 + # Verify that the expected number of complete messages has arrived + assert expected_msgs_received_count == len(self.transfer.get_messages_in()) + + # Get the data as received by DelugeTransferProtocol + message1 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg1) == rencode.dumps(message1) + message2 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg2) == rencode.dumps(message2) + message3 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg1) == rencode.dumps(message3) + + # Remove underscore to enable test, or run the test directly: + def _test_rencode_fail_protocol(self): + """ + This test tries to test the protocol that relies on errors from rencode. + + """ + msg_bytes = ( + base64.b64decode(self.msg1_expected_compressed_base64) + + base64.b64decode(self.msg2_expected_compressed_base64) + + base64.b64decode(self.msg1_expected_compressed_base64) + ) + packet_size = 149 + + one_message_byte_count = len( + base64.b64decode(self.msg1_expected_compressed_base64) + ) + two_messages_byte_count = one_message_byte_count + len( + base64.b64decode(self.msg2_expected_compressed_base64) + ) + three_messages_byte_count = two_messages_byte_count + len( + base64.b64decode(self.msg1_expected_compressed_base64) + ) + + print() + + print( + 'Msg1 size:', + len(base64.b64decode(self.msg1_expected_compressed_base64)) - 4, + ) + print( + 'Msg2 size:', + len(base64.b64decode(self.msg2_expected_compressed_base64)) - 4, + ) + print( + 'Msg3 size:', + len(base64.b64decode(self.msg1_expected_compressed_base64)) - 4, + ) + + print('one_message_byte_count:', one_message_byte_count) + print('two_messages_byte_count:', two_messages_byte_count) + print('three_messages_byte_count:', three_messages_byte_count) + + for d in self.receive_parts_helper( + msg_bytes, packet_size, self.transfer.data_received_old_protocol + ): + bytes_received = self.transfer.get_bytes_recv() + + if bytes_received >= three_messages_byte_count: + expected_msgs_received_count = 3 + elif bytes_received >= two_messages_byte_count: + expected_msgs_received_count = 2 + elif bytes_received >= one_message_byte_count: + expected_msgs_received_count = 1 + else: + expected_msgs_received_count = 0 + # Verify that the expected number of complete messages has arrived + if expected_msgs_received_count != len(self.transfer.get_messages_in()): + print( + 'Expected number of messages received is %d, but %d have been received.' + % ( + expected_msgs_received_count, + len(self.transfer.get_messages_in()), + ) + ) + + # Get the data as received by DelugeTransferProtocol + message1 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg1) == rencode.dumps(message1) + message2 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg2) == rencode.dumps(message2) + message3 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg1) == rencode.dumps(message3) + + def test_receive_middle_of_header(self): + """ + This test concatenates two messsages (as they're sent over the network), + and lets DelugeTransferProtocol receive the data in two parts. + The first part contains the first message, plus two bytes of the next message. + The next part contains the rest of the message. + + This is a special case, as DelugeTransferProtocol can't start parsing + a message until it has at least 5 bytes (the size of the header) to be able + to read and parse the size of the payload. + + """ + two_concatenated = base64.b64decode( + self.msg1_expected_compressed_base64 + ) + base64.b64decode(self.msg2_expected_compressed_base64) + first_len = len(base64.b64decode(self.msg1_expected_compressed_base64)) + + # Now found the entire first message, and half the header of the next message (2 bytes into the header) + self.transfer.dataReceived(two_concatenated[: first_len + 2]) + + # Should be 1 message in the list + assert 1 == len(self.transfer.get_messages_in()) + + # Send the rest + self.transfer.dataReceived(two_concatenated[first_len + 2 :]) + + # Should be 2 messages in the list + assert 2 == len(self.transfer.get_messages_in()) + + # Get the data as sent by DelugeTransferProtocol + message1 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg1) == rencode.dumps(message1) + message2 = self.transfer.get_messages_in().pop(0) + assert rencode.dumps(self.msg2) == rencode.dumps(message2) + + # Needs file containing big data structure e.g. like thetorrent list as it is transfered by the daemon + # def test_simulate_big_transfer(self): + # filename = '../deluge.torrentlist' + # + # f = open(filename, 'r') + # data = f.read() + # message_to_send = eval(data) + # self.transfer.transfer_message(message_to_send) + # + # Get the data as sent to the network by DelugeTransferProtocol + # compressed_data = self.transfer.get_messages_out_joined() + # packet_size = 16000 # Or something smaller... + # + # for d in self.receive_parts_helper(compressed_data, packet_size): + # bytes_recv = self.transfer.get_bytes_recv() + # if bytes_recv < len(compressed_data): + # self.assertEqual(len(self.transfer.get_messages_in()), 0) + # else: + # self.assertEqual(len(self.transfer.get_messages_in()), 1) + # Get the data as received by DelugeTransferProtocol + # transfered_message = self.transfer.get_messages_in().pop(0) + # Test that the data structures are equal + # self.assertEqual(transfered_message, message_to_send) + # self.assertTrue(transfered_message == message_to_send) + # + # f.close() + # f = open('rencode.torrentlist', 'w') + # f.write(str(transfered_message)) + # f.close() + + def receive_parts_helper(self, data, packet_size, receive_func=None): + byte_count = len(data) + sent_bytes = 0 + while byte_count > 0: + to_receive = packet_size if byte_count > packet_size else byte_count + sent_bytes += to_receive + byte_count -= to_receive + if receive_func: + receive_func(data[:to_receive]) + else: + self.transfer.dataReceived(data[:to_receive]) + data = data[to_receive:] + yield |