Deluge
diff options
Diffstat (limited to 'tests/test_transfer.py')
| -rw-r--r-- | tests/test_transfer.py | 368 |
1 files changed, 368 insertions, 0 deletions
diff --git a/tests/test_transfer.py b/tests/test_transfer.py new file mode 100644 index 000000000..cd9a271d5 --- /dev/null +++ b/tests/test_transfer.py @@ -0,0 +1,368 @@ +# -*- coding: utf-8 -*- +# +# test_transfer.py +# +# Copyright (C) 2012 Bro <bro.development@gmail.com> +# +# Deluge is free software. +# +# You may redistribute it and/or modify it under the terms of the +# GNU General Public License, as published by the Free Software +# Foundation; either version 3 of the License, or (at your option) +# any later version. +# +# deluge is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. +# See the GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with deluge. If not, write to: +# The Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor +# Boston, MA 02110-1301, USA. +# +# In addition, as a special exception, the copyright holders give +# permission to link the code of portions of this program with the OpenSSL +# library. +# You must obey the GNU General Public License in all respects for all of +# the code used other than OpenSSL. If you modify file(s) with this +# exception, you may extend this exception to your version of the file(s), +# but you are not obligated to do so. If you do not wish to do so, delete +# this exception statement from your version. If you delete this exception +# statement from all source files in the program, then also delete it here. +# + +from twisted.trial import unittest + +from deluge.transfer import DelugeTransferProtocol + +import base64 + +import deluge.rencode as rencode + +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 dataReceived_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. + + """ + from datetime import timedelta + 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", + 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, e: + #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(e)) + return + else: + data = dobj.unused_data + self._message_length = 0 + + self.message_received(request) + +class DelugeTransferProtocolTestCase(unittest.TestCase): + + def setUp(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 pickle, + 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": "some string"}, {"key_bool": True}) + self.msg2 = (2, 3, {"key_float": 12424.29423}, + {"key_unicode": u"some string"}, + {"key_dict_with_tuple": {"key_tuple": (1, 2, 3)}}, + {"keylist": [4, "5", 6.7]}) + + self.msg1_expected_compressed_base64 = "RAAAADF4nDvKwJjenp1aGZ+ZV+Lgxfv9PYRXXFLU"\ + "XZyfm6oAZGTmpad3gAST8vNznAEAJhSQ" + + self.msg2_expected_compressed_base64 = "RAAAAF14nDvGxJzemZ1aGZ+Wk59Y4uTmpKib3g3i"\ + "l+ZlJuenpHYX5+emKhSXFGXmpadPBkmkZCaXxJdn"\ + "lmTEl5QW5KRCdIOZhxmBhrUDuTmZxSWHWRpNnRyu"\ + "paUBAHYlJxI=" + + 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) + self.assertEquals(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) + self.assertEquals(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)) + self.assertEquals(len(self.transfer.get_messages_in()), 0) + self.assertEquals(self.transfer._message_length, 0) + self.assertEquals(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) + self.assertEquals(rencode.dumps(self.msg1), rencode.dumps(message1)) + message2 = self.transfer.get_messages_in().pop(0) + self.assertEquals(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 + self.assertEquals(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) + self.assertEquals(rencode.dumps(self.msg1), rencode.dumps(message1)) + message2 = self.transfer.get_messages_in().pop(0) + self.assertEquals(rencode.dumps(self.msg2), rencode.dumps(message2)) + message3 = self.transfer.get_messages_in().pop(0) + self.assertEquals(rencode.dumps(self.msg1), rencode.dumps(message3)) + + + # Remove underscore to enable test, or run the test directly: + # tests $ trial test_transfer.DelugeTransferProtocolTestCase._test_rencode_fail_protocol + 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.dataReceived_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) + self.assertEquals(rencode.dumps(self.msg1), rencode.dumps(message1)) + message2 = self.transfer.get_messages_in().pop(0) + self.assertEquals(rencode.dumps(self.msg2), rencode.dumps(message2)) + message3 = self.transfer.get_messages_in().pop(0) + self.assertEquals(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 4 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 + self.assertEquals(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 + self.assertEquals(2, len(self.transfer.get_messages_in())) + + # Get the data as sent by DelugeTransferProtocol + message1 = self.transfer.get_messages_in().pop(0) + self.assertEquals(rencode.dumps(self.msg1), rencode.dumps(message1)) + message2 = self.transfer.get_messages_in().pop(0) + self.assertEquals(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.assertEquals(len(self.transfer.get_messages_in()), 0) + # else: + # self.assertEquals(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.assertEquals(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 |