OLD | NEW |
| (Empty) |
1 # Copyright (c) 2001-2007 Twisted Matrix Laboratories. | |
2 # See LICENSE for details. | |
3 | |
4 | |
5 """ | |
6 Test cases for twisted.protocols.stateful | |
7 """ | |
8 | |
9 from twisted.test import test_protocols | |
10 from twisted.protocols.stateful import StatefulProtocol | |
11 | |
12 from struct import pack, unpack, calcsize | |
13 | |
14 | |
15 class MyInt32StringReceiver(StatefulProtocol): | |
16 """ | |
17 A stateful Int32StringReceiver. | |
18 """ | |
19 MAX_LENGTH = 99999 | |
20 structFormat = "!I" | |
21 prefixLength = calcsize(structFormat) | |
22 | |
23 def getInitialState(self): | |
24 return self._getHeader, 4 | |
25 | |
26 def _getHeader(self, msg): | |
27 length, = unpack("!i", msg) | |
28 if length > self.MAX_LENGTH: | |
29 self.transport.loseConnection() | |
30 return | |
31 return self._getString, length | |
32 | |
33 def _getString(self, msg): | |
34 self.stringReceived(msg) | |
35 return self._getHeader, 4 | |
36 | |
37 def stringReceived(self, msg): | |
38 """ | |
39 Override this. | |
40 """ | |
41 raise NotImplementedError | |
42 | |
43 def sendString(self, data): | |
44 """ | |
45 Send an int32-prefixed string to the other end of the connection. | |
46 """ | |
47 self.transport.write(pack(self.structFormat, len(data)) + data) | |
48 | |
49 | |
50 class TestInt32(MyInt32StringReceiver): | |
51 def connectionMade(self): | |
52 self.received = [] | |
53 | |
54 def stringReceived(self, s): | |
55 self.received.append(s) | |
56 | |
57 MAX_LENGTH = 50 | |
58 closed = 0 | |
59 | |
60 def connectionLost(self, reason): | |
61 self.closed = 1 | |
62 | |
63 | |
64 class Int32TestCase(test_protocols.Int32TestCase): | |
65 protocol = TestInt32 | |
66 | |
67 def test_bigReceive(self): | |
68 r = self.getProtocol() | |
69 big = "" | |
70 for s in self.strings * 4: | |
71 big += pack("!i", len(s)) + s | |
72 r.dataReceived(big) | |
73 self.assertEquals(r.received, self.strings * 4) | |
74 | |
OLD | NEW |