summaryrefslogtreecommitdiffstats
path: root/tests/test_transfer.py
blob: cd9a271d584c8686bc325babc91ca3ff5e785f90 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
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