BtPortMessageTest.cc 5.12 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
#include "BtPortMessage.h"

#include <cstring>

#include <cppunit/extensions/HelperMacros.h>

#include "bittorrent_helper.h"
#include "util.h"
#include "array_fun.h"
#include "Peer.h"
#include "DHTNode.h"
#include "DHTRoutingTable.h"
#include "MockDHTTask.h"
#include "MockDHTTaskFactory.h"
#include "MockDHTTaskQueue.h"

namespace aria2 {

19
class BtPortMessageTest : public CppUnit::TestFixture {
20 21 22 23 24 25 26 27 28

  CPPUNIT_TEST_SUITE(BtPortMessageTest);
  CPPUNIT_TEST(testCreate);
  CPPUNIT_TEST(testToString);
  CPPUNIT_TEST(testCreateMessage);
  CPPUNIT_TEST(testDoReceivedAction);
  CPPUNIT_TEST(testDoReceivedAction_bootstrap);
  CPPUNIT_TEST_SUITE_END();

29
private:
30
public:
31
  void setUp() {}
32 33 34 35 36 37 38

  void testCreate();
  void testToString();
  void testCreateMessage();
  void testDoReceivedAction();
  void testDoReceivedAction_bootstrap();

39
  class MockDHTTaskFactory2 : public MockDHTTaskFactory {
40
  public:
41
    virtual std::shared_ptr<DHTTask>
42 43
    createPingTask(const std::shared_ptr<DHTNode>& remoteNode,
                   int numRetry) CXX11_OVERRIDE
44
    {
45
      return std::shared_ptr<DHTTask>(new MockDHTTask(remoteNode));
46 47
    }

48 49
    virtual std::shared_ptr<DHTTask>
    createNodeLookupTask(const unsigned char* targetID) CXX11_OVERRIDE
50
    {
51 52
      std::shared_ptr<MockDHTTask> task(
          new MockDHTTask(std::shared_ptr<DHTNode>()));
53 54 55 56 57 58 59 60
      task->setTargetID(targetID);
      return task;
    }
  };
};

CPPUNIT_TEST_SUITE_REGISTRATION(BtPortMessageTest);

61 62
void BtPortMessageTest::testCreate()
{
63 64 65
  unsigned char msg[7];
  bittorrent::createPeerMessageString(msg, sizeof(msg), 3, 9);
  bittorrent::setShortIntParam(&msg[5], 12345);
66
  std::shared_ptr<BtPortMessage> pm(BtPortMessage::create(&msg[4], 3));
67 68 69 70 71 72 73 74 75
  CPPUNIT_ASSERT_EQUAL((uint8_t)9, pm->getId());
  CPPUNIT_ASSERT_EQUAL((uint16_t)12345, pm->getPort());

  // case: payload size is wrong
  try {
    unsigned char msg[8];
    bittorrent::createPeerMessageString(msg, sizeof(msg), 4, 9);
    BtPortMessage::create(&msg[4], 4);
    CPPUNIT_FAIL("exception must be thrown.");
76 77
  }
  catch (...) {
78 79 80 81 82 83 84
  }
  // case: id is wrong
  try {
    unsigned char msg[7];
    bittorrent::createPeerMessageString(msg, sizeof(msg), 3, 10);
    BtPortMessage::create(&msg[4], 3);
    CPPUNIT_FAIL("exception must be thrown.");
85 86
  }
  catch (...) {
87 88 89
  }
}

90 91
void BtPortMessageTest::testToString()
{
92 93 94 95
  BtPortMessage msg(1);
  CPPUNIT_ASSERT_EQUAL(std::string("port port=1"), msg.toString());
}

96 97
void BtPortMessageTest::testCreateMessage()
{
98 99 100 101
  BtPortMessage msg(6881);
  unsigned char data[7];
  bittorrent::createPeerMessageString(data, sizeof(data), 3, 9);
  bittorrent::setShortIntParam(&data[5], 6881);
102 103 104
  auto rawmsg = msg.createMessage();
  CPPUNIT_ASSERT_EQUAL((size_t)7, rawmsg.size());
  CPPUNIT_ASSERT(std::equal(std::begin(rawmsg), std::end(rawmsg), data));
105 106 107 108 109 110
}

void BtPortMessageTest::testDoReceivedAction()
{
  unsigned char nodeID[DHT_ID_LENGTH];
  memset(nodeID, 0, DHT_ID_LENGTH);
111
  std::shared_ptr<DHTNode> localNode(new DHTNode(nodeID));
112 113

  // 9 nodes to create at least 2 buckets.
114
  std::shared_ptr<DHTNode> nodes[9];
115
  for (size_t i = 0; i < arraySize(nodes); ++i) {
116
    memset(nodeID, 0, DHT_ID_LENGTH);
117
    nodeID[DHT_ID_LENGTH - 1] = i + 1;
118 119 120 121
    nodes[i].reset(new DHTNode(nodeID));
  }

  DHTRoutingTable routingTable(localNode);
122
  for (size_t i = 0; i < arraySize(nodes); ++i) {
123 124 125
    routingTable.addNode(nodes[i]);
  }

126
  std::shared_ptr<Peer> peer(new Peer("192.168.0.1", 6881));
127 128 129 130 131 132 133 134 135 136 137 138 139
  BtPortMessage msg(6881);
  MockDHTTaskQueue taskQueue;
  MockDHTTaskFactory2 taskFactory;
  msg.setLocalNode(localNode.get());
  msg.setRoutingTable(&routingTable);
  msg.setTaskQueue(&taskQueue);
  msg.setTaskFactory(&taskFactory);
  msg.setPeer(peer);

  msg.doReceivedAction();

  CPPUNIT_ASSERT_EQUAL((size_t)1, taskQueue.immediateTaskQueue_.size());

140 141
  auto task =
      std::dynamic_pointer_cast<MockDHTTask>(taskQueue.immediateTaskQueue_[0]);
142
  std::shared_ptr<DHTNode> node = task->remoteNode_;
143 144 145 146 147 148 149 150 151
  CPPUNIT_ASSERT_EQUAL(std::string("192.168.0.1"), node->getIPAddress());
  CPPUNIT_ASSERT_EQUAL((uint16_t)6881, node->getPort());
}

void BtPortMessageTest::testDoReceivedAction_bootstrap()
{
  unsigned char nodeID[DHT_ID_LENGTH];
  memset(nodeID, 0, DHT_ID_LENGTH);
  nodeID[0] = 0xff;
152
  std::shared_ptr<DHTNode> localNode(new DHTNode(nodeID));
153 154
  DHTRoutingTable routingTable(localNode); // no nodes , 1 bucket.

155
  std::shared_ptr<Peer> peer(new Peer("192.168.0.1", 6881));
156 157 158 159 160 161 162 163 164 165 166 167
  BtPortMessage msg(6881);
  MockDHTTaskQueue taskQueue;
  MockDHTTaskFactory2 taskFactory;
  msg.setLocalNode(localNode.get());
  msg.setRoutingTable(&routingTable);
  msg.setTaskQueue(&taskQueue);
  msg.setTaskFactory(&taskFactory);
  msg.setPeer(peer);

  msg.doReceivedAction();

  CPPUNIT_ASSERT_EQUAL((size_t)2, taskQueue.immediateTaskQueue_.size());
168 169
  auto task =
      std::dynamic_pointer_cast<MockDHTTask>(taskQueue.immediateTaskQueue_[0]);
170
  std::shared_ptr<DHTNode> node(task->remoteNode_);
171 172 173
  CPPUNIT_ASSERT_EQUAL(std::string("192.168.0.1"), node->getIPAddress());
  CPPUNIT_ASSERT_EQUAL((uint16_t)6881, node->getPort());

174 175
  auto task2 =
      std::dynamic_pointer_cast<MockDHTTask>(taskQueue.immediateTaskQueue_[1]);
176 177 178 179
  CPPUNIT_ASSERT(memcmp(nodeID, task2->targetID_, DHT_ID_LENGTH) == 0);
}

} // namespace aria2