DefaultBtRequestFactoryTest.cc 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244
  1. #include "DefaultBtRequestFactory.h"
  2. #include <algorithm>
  3. #include <cppunit/extensions/HelperMacros.h>
  4. #include "MockBtMessage.h"
  5. #include "MockBtMessageFactory.h"
  6. #include "MockBtMessageDispatcher.h"
  7. #include "MockPieceStorage.h"
  8. #include "Peer.h"
  9. #include "FileEntry.h"
  10. #include "BtHandshakeMessage.h"
  11. #include "DownloadContext.h"
  12. #include "bittorrent_helper.h"
  13. namespace aria2 {
  14. class DefaultBtRequestFactoryTest:public CppUnit::TestFixture {
  15. CPPUNIT_TEST_SUITE(DefaultBtRequestFactoryTest);
  16. CPPUNIT_TEST(testAddTargetPiece);
  17. CPPUNIT_TEST(testRemoveCompletedPiece);
  18. CPPUNIT_TEST(testCreateRequestMessages);
  19. CPPUNIT_TEST(testCreateRequestMessages_onEndGame);
  20. CPPUNIT_TEST(testRemoveTargetPiece);
  21. CPPUNIT_TEST(testGetTargetPieceIndexes);
  22. CPPUNIT_TEST_SUITE_END();
  23. private:
  24. SharedHandle<Peer> _peer;
  25. SharedHandle<DefaultBtRequestFactory> _requestFactory;
  26. SharedHandle<DownloadContext> _dctx;
  27. SharedHandle<MockPieceStorage> _pieceStorage;
  28. SharedHandle<MockBtMessageFactory> _messageFactory;
  29. SharedHandle<MockBtMessageDispatcher> _dispatcher;
  30. public:
  31. void testAddTargetPiece();
  32. void testRemoveCompletedPiece();
  33. void testCreateRequestMessages();
  34. void testCreateRequestMessages_onEndGame();
  35. void testRemoveTargetPiece();
  36. void testGetTargetPieceIndexes();
  37. class MockBtRequestMessage : public MockBtMessage {
  38. public:
  39. size_t index;
  40. size_t blockIndex;
  41. MockBtRequestMessage(size_t index, size_t blockIndex):
  42. index(index), blockIndex(blockIndex) {}
  43. };
  44. typedef SharedHandle<MockBtRequestMessage> MockBtRequestMessageHandle;
  45. class MockBtMessageFactory2 : public MockBtMessageFactory {
  46. public:
  47. virtual SharedHandle<BtMessage>
  48. createRequestMessage(const SharedHandle<Piece>& piece, size_t blockIndex) {
  49. return SharedHandle<BtMessage>
  50. (new MockBtRequestMessage(piece->getIndex(), blockIndex));
  51. }
  52. };
  53. class MockBtMessageDispatcher2 : public MockBtMessageDispatcher {
  54. public:
  55. virtual bool isOutstandingRequest(size_t index, size_t blockIndex) {
  56. return index == 0 && blockIndex == 0;
  57. }
  58. };
  59. class SortMockBtRequestMessage {
  60. public:
  61. bool operator()(const SharedHandle<MockBtRequestMessage>& a,
  62. const SharedHandle<MockBtRequestMessage>& b) {
  63. if(a->index < b->index) {
  64. return true;
  65. } else if(b->index < a->index) {
  66. return false;
  67. } else if(a->blockIndex < b->blockIndex) {
  68. return true;
  69. } else if(b->blockIndex < a->blockIndex) {
  70. return false;
  71. } else {
  72. return true;
  73. }
  74. }
  75. };
  76. void setUp()
  77. {
  78. _pieceStorage.reset(new MockPieceStorage());
  79. _peer.reset(new Peer("host", 6969));
  80. _messageFactory.reset(new MockBtMessageFactory2());
  81. _dispatcher.reset(new MockBtMessageDispatcher());
  82. _requestFactory.reset(new DefaultBtRequestFactory());
  83. _requestFactory->setPieceStorage(_pieceStorage);
  84. _requestFactory->setPeer(_peer);
  85. _requestFactory->setBtMessageDispatcher(_dispatcher);
  86. _requestFactory->setBtMessageFactory(_messageFactory);
  87. }
  88. };
  89. CPPUNIT_TEST_SUITE_REGISTRATION(DefaultBtRequestFactoryTest);
  90. void DefaultBtRequestFactoryTest::testAddTargetPiece() {
  91. {
  92. SharedHandle<Piece> piece(new Piece(0, 16*1024*10));
  93. _requestFactory->addTargetPiece(piece);
  94. CPPUNIT_ASSERT_EQUAL((size_t)1, _requestFactory->countTargetPiece());
  95. }
  96. {
  97. SharedHandle<Piece> piece(new Piece(1, 16*1024*9));
  98. piece->completeBlock(0);
  99. _requestFactory->addTargetPiece(piece);
  100. CPPUNIT_ASSERT_EQUAL((size_t)2, _requestFactory->countTargetPiece());
  101. }
  102. CPPUNIT_ASSERT_EQUAL((size_t)18, _requestFactory->countMissingBlock());
  103. }
  104. void DefaultBtRequestFactoryTest::testRemoveCompletedPiece() {
  105. SharedHandle<Piece> piece1(new Piece(0, 16*1024));
  106. SharedHandle<Piece> piece2(new Piece(1, 16*1024));
  107. piece2->setAllBlock();
  108. _requestFactory->addTargetPiece(piece1);
  109. _requestFactory->addTargetPiece(piece2);
  110. CPPUNIT_ASSERT_EQUAL((size_t)2, _requestFactory->countTargetPiece());
  111. _requestFactory->removeCompletedPiece();
  112. CPPUNIT_ASSERT_EQUAL((size_t)1, _requestFactory->countTargetPiece());
  113. CPPUNIT_ASSERT_EQUAL((size_t)0,
  114. _requestFactory->getTargetPieces().front()->getIndex());
  115. }
  116. void DefaultBtRequestFactoryTest::testCreateRequestMessages() {
  117. int PIECE_LENGTH = 16*1024*2;
  118. SharedHandle<Piece> piece1(new Piece(0, PIECE_LENGTH));
  119. SharedHandle<Piece> piece2(new Piece(1, PIECE_LENGTH));
  120. _requestFactory->addTargetPiece(piece1);
  121. _requestFactory->addTargetPiece(piece2);
  122. std::vector<SharedHandle<BtMessage> > msgs;
  123. _requestFactory->createRequestMessages(msgs, 3);
  124. CPPUNIT_ASSERT_EQUAL((size_t)3, msgs.size());
  125. std::vector<SharedHandle<BtMessage> >::iterator itr = msgs.begin();
  126. SharedHandle<MockBtRequestMessage> msg =
  127. dynamic_pointer_cast<MockBtRequestMessage>(*itr);
  128. CPPUNIT_ASSERT_EQUAL((size_t)0, msg->index);
  129. CPPUNIT_ASSERT_EQUAL((size_t)0, msg->blockIndex);
  130. ++itr;
  131. msg = dynamic_pointer_cast<MockBtRequestMessage>(*itr);
  132. CPPUNIT_ASSERT_EQUAL((size_t)0, msg->index);
  133. CPPUNIT_ASSERT_EQUAL((size_t)1, msg->blockIndex);
  134. ++itr;
  135. msg = dynamic_pointer_cast<MockBtRequestMessage>(*itr);
  136. CPPUNIT_ASSERT_EQUAL((size_t)1, msg->index);
  137. CPPUNIT_ASSERT_EQUAL((size_t)0, msg->blockIndex);
  138. {
  139. std::vector<SharedHandle<BtMessage> > msgs;
  140. _requestFactory->createRequestMessages(msgs, 3);
  141. CPPUNIT_ASSERT_EQUAL((size_t)1, msgs.size());
  142. }
  143. }
  144. void DefaultBtRequestFactoryTest::testCreateRequestMessages_onEndGame() {
  145. SharedHandle<MockBtMessageDispatcher2> dispatcher
  146. (new MockBtMessageDispatcher2());
  147. _requestFactory->setBtMessageDispatcher(dispatcher);
  148. int PIECE_LENGTH = 16*1024*2;
  149. SharedHandle<Piece> piece1(new Piece(0, PIECE_LENGTH));
  150. SharedHandle<Piece> piece2(new Piece(1, PIECE_LENGTH));
  151. _requestFactory->addTargetPiece(piece1);
  152. _requestFactory->addTargetPiece(piece2);
  153. std::vector<SharedHandle<BtMessage> > msgs;
  154. _requestFactory->createRequestMessagesOnEndGame(msgs, 3);
  155. std::vector<SharedHandle<MockBtRequestMessage> > mmsgs;
  156. for(std::vector<SharedHandle<BtMessage> >::iterator i = msgs.begin();
  157. i != msgs.end(); ++i) {
  158. mmsgs.push_back(dynamic_pointer_cast<MockBtRequestMessage>(*i));
  159. }
  160. std::sort(mmsgs.begin(), mmsgs.end(), SortMockBtRequestMessage());
  161. CPPUNIT_ASSERT_EQUAL((size_t)3, mmsgs.size());
  162. std::vector<SharedHandle<MockBtRequestMessage> >::iterator itr =mmsgs.begin();
  163. MockBtRequestMessage* msg = (*itr).get();
  164. CPPUNIT_ASSERT_EQUAL((size_t)0, msg->index);
  165. CPPUNIT_ASSERT_EQUAL((size_t)1, msg->blockIndex);
  166. ++itr;
  167. msg = (*itr).get();
  168. CPPUNIT_ASSERT_EQUAL((size_t)1, msg->index);
  169. CPPUNIT_ASSERT_EQUAL((size_t)0, msg->blockIndex);
  170. ++itr;
  171. msg = (*itr).get();
  172. CPPUNIT_ASSERT_EQUAL((size_t)1, msg->index);
  173. CPPUNIT_ASSERT_EQUAL((size_t)1, msg->blockIndex);
  174. }
  175. void DefaultBtRequestFactoryTest::testRemoveTargetPiece() {
  176. SharedHandle<Piece> piece1(new Piece(0, 16*1024));
  177. _requestFactory->addTargetPiece(piece1);
  178. CPPUNIT_ASSERT(std::find(_requestFactory->getTargetPieces().begin(),
  179. _requestFactory->getTargetPieces().end(),
  180. piece1) !=
  181. _requestFactory->getTargetPieces().end());
  182. _requestFactory->removeTargetPiece(piece1);
  183. CPPUNIT_ASSERT(std::find(_requestFactory->getTargetPieces().begin(),
  184. _requestFactory->getTargetPieces().end(),
  185. piece1) ==
  186. _requestFactory->getTargetPieces().end());
  187. }
  188. void DefaultBtRequestFactoryTest::testGetTargetPieceIndexes()
  189. {
  190. SharedHandle<Piece> piece1(new Piece(1, 16*1024));
  191. SharedHandle<Piece> piece3(new Piece(3, 16*1024));
  192. SharedHandle<Piece> piece5(new Piece(5, 16*1024));
  193. _requestFactory->addTargetPiece(piece3);
  194. _requestFactory->addTargetPiece(piece1);
  195. _requestFactory->addTargetPiece(piece5);
  196. std::vector<size_t> indexes;
  197. _requestFactory->getTargetPieceIndexes(indexes);
  198. CPPUNIT_ASSERT_EQUAL((size_t)3, indexes.size());
  199. CPPUNIT_ASSERT_EQUAL((size_t)3, indexes[0]);
  200. CPPUNIT_ASSERT_EQUAL((size_t)1, indexes[1]);
  201. CPPUNIT_ASSERT_EQUAL((size_t)5, indexes[2]);
  202. }
  203. } // namespace aria2