1
0

test_krpc.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.1 (the License). You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License. You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10. from unittest import *
  11. from krpc import *
  12. from BitTorrent.defaultargs import common_options, rare_options
  13. from BTL.stackthreading import Event
  14. from BitTorrent.RawServer_twisted import RawServer
  15. from node import Node
  16. KRPC.noisy = 0
  17. import sys
  18. if __name__ =="__main__":
  19. tests = defaultTestLoader.loadTestsFromNames([sys.argv[0][:-3]])
  20. result = TextTestRunner().run(tests)
  21. def connectionForAddr(host, port):
  22. return host
  23. class Receiver(object):
  24. protocol = KRPC
  25. def __init__(self, addr):
  26. self.buf = []
  27. self.node = Node().init('0'*20, addr[0], addr[1])
  28. def krpc_store(self, msg, _krpc_sender):
  29. self.buf += [msg]
  30. def krpc_echo(self, msg, _krpc_sender):
  31. return msg
  32. class KRPCTests(TestCase):
  33. def setUp(self):
  34. self.noisy = 0
  35. d = dict([(x[0],x[1]) for x in common_options + rare_options])
  36. self.r = RawServer(d)
  37. addr = ('127.0.0.1', 1180)
  38. self.as = self.r.create_udpsocket(addr[1], addr[0], True)
  39. self.af = Receiver(addr)
  40. self.a = hostbroker(self.af, addr, self.as, self.r.add_task)
  41. self.r.start_listening_udp(self.as, self.a)
  42. addr = ('127.0.0.1', 1181)
  43. self.bs = self.r.create_udpsocket(addr[1], addr[0], True)
  44. self.bf = Receiver(addr)
  45. self.b = hostbroker(self.bf, addr, self.bs, self.r.add_task)
  46. self.r.start_listening_udp(self.bs, self.b)
  47. def tearDown(self):
  48. self.as.close()
  49. self.bs.close()
  50. def testSimpleMessage(self):
  51. self.noisy = 0
  52. self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('store', {'msg' : "This is a test."})
  53. self.r.listen_once(0.01)
  54. self.assertEqual(self.bf.buf, ["This is a test."])
  55. def testMessageBlast(self):
  56. self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('store', {'msg' : "This is a test."})
  57. self.r.listen_once(0.01)
  58. self.assertEqual(self.bf.buf, ["This is a test."])
  59. self.bf.buf = []
  60. for i in range(100):
  61. self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('store', {'msg' : "This is a test."})
  62. self.r.listen_once(0.01)
  63. #self.bf.buf = []
  64. self.assertEqual(self.bf.buf, ["This is a test."] * 100)
  65. def testEcho(self):
  66. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
  67. df.addCallback(self.gotMsg)
  68. self.r.listen_once(0.01)
  69. self.r.listen_once(0.01)
  70. self.assertEqual(self.msg, "This is a test.")
  71. def gotMsg(self, dict):
  72. _krpc_sender = dict['_krpc_sender']
  73. msg = dict['rsp']
  74. self.msg = msg
  75. def testManyEcho(self):
  76. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
  77. df.addCallback(self.gotMsg)
  78. self.r.listen_once(0.01)
  79. self.r.listen_once(0.01)
  80. self.assertEqual(self.msg, "This is a test.")
  81. for i in xrange(100):
  82. self.msg = None
  83. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
  84. df.addCallback(self.gotMsg)
  85. self.r.listen_once(0.01)
  86. self.r.listen_once(0.01)
  87. self.assertEqual(self.msg, "This is a test.")
  88. def testMultiEcho(self):
  89. self.noisy = 0
  90. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
  91. df.addCallback(self.gotMsg)
  92. self.r.listen_once(0.01)
  93. self.r.listen_once(0.01)
  94. self.assertEqual(self.msg, "This is a test.")
  95. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is another test."})
  96. df.addCallback(self.gotMsg)
  97. self.r.listen_once(0.01)
  98. self.r.listen_once(0.01)
  99. self.assertEqual(self.msg, "This is another test.")
  100. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is yet another test."})
  101. df.addCallback(self.gotMsg)
  102. self.r.listen_once(0.01)
  103. self.r.listen_once(0.01)
  104. self.assertEqual(self.msg, "This is yet another test.")
  105. def testEchoReset(self):
  106. self.noisy = 0
  107. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
  108. df.addCallback(self.gotMsg)
  109. self.r.listen_once(0.01)
  110. self.r.listen_once(0.01)
  111. self.assertEqual(self.msg, "This is a test.")
  112. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is another test."})
  113. df.addCallback(self.gotMsg)
  114. self.r.listen_once(0.01)
  115. self.r.listen_once(0.01)
  116. self.assertEqual(self.msg, "This is another test.")
  117. del(self.a.connections[('127.0.0.1', 1181)])
  118. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is yet another test."})
  119. df.addCallback(self.gotMsg)
  120. self.r.listen_once(0.01)
  121. self.r.listen_once(0.01)
  122. self.assertEqual(self.msg, "This is yet another test.")
  123. def testLotsofEchoReset(self):
  124. for i in range(100):
  125. self.testEchoReset()
  126. def testUnknownMeth(self):
  127. self.noisy = 0
  128. df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('blahblah', {'msg' : "This is a test."})
  129. df.addErrback(self.gotErr)
  130. self.r.listen_once(0.01)
  131. self.r.listen_once(0.01)
  132. self.assertEqual(self.err, KRPC_ERROR_METHOD_UNKNOWN)
  133. def gotErr(self, err):
  134. self.err = err