2009-11-23 Tatsuhiro Tsujikawa <t-tujikawa@users.sourceforge.net>

	Fixed ut_metadata data handling. Implemented
	UTMetadataDataExtensionMessage::doReceivedAction().  Initialize
	PeerStorage in HandshakeExtensionMessage::doReceivedAction() when
	metadata_size is received.
	* src/DefaultExtensionMessageFactory.cc
	* src/DefaultExtensionMessageFactory.h
	* src/HandshakeExtensionMessage.cc
	* src/HandshakeExtensionMessage.h
	* src/UTMetadataDataExtensionMessage.cc
	* src/UTMetadataDataExtensionMessage.h
	* src/UTMetadataRequestExtensionMessage.cc
	* src/UTMetadataRequestFactory.cc
	* src/UTMetadataRequestFactory.h
	* src/UTMetadataRequestTracker.cc
	* src/UTMetadataRequestTracker.h
	* test/DefaultExtensionMessageFactoryTest.cc
	* test/HandshakeExtensionMessageTest.cc
	* test/MockBtMessage.h
	* test/UTMetadataDataExtensionMessageTest.cc
	* test/UTMetadataRequestExtensionMessageTest.cc
	* test/UTMetadataRequestFactoryTest.cc
	* test/UTMetadataRequestTrackerTest.cc
	* test/extension_message_test_helper.h

8 files changed, 269 insertions, 25 deletions

diff --git a/test/DefaultExtensionMessageFactoryTest.cc b/test/DefaultExtensionMessageFactoryTest.cc
index 878bdb06..9ace3e29 100644
--- a/test/DefaultExtensionMessageFactoryTest.cc
+++ b/test/DefaultExtensionMessageFactoryTest.cc
@@ -21,6 +21,10 @@
 #include "UTMetadataRequestExtensionMessage.h"
 #include "UTMetadataDataExtensionMessage.h"
 #include "UTMetadataRejectExtensionMessage.h"
+#include "BtRuntime.h"
+#include "PieceStorage.h"
+#include "RequestGroup.h"
+#include "Option.h"
 
 namespace aria2 {
 
@@ -42,6 +46,7 @@ private:
   SharedHandle<MockBtMessageDispatcher> _dispatcher;
   SharedHandle<MockBtMessageFactory> _messageFactory;
   SharedHandle<DownloadContext> _dctx;
+  SharedHandle<RequestGroup> _requestGroup;
 public:
   void setUp()
   {
@@ -59,6 +64,11 @@ public:
 
     _dctx.reset(new DownloadContext());
 
+    SharedHandle<Option> option(new Option());
+    _requestGroup.reset(new RequestGroup(option));
+    _requestGroup->setDownloadContext(_dctx);
+    _dctx->setOwnerRequestGroup(_requestGroup.get());
+
     _factory.reset(new DefaultExtensionMessageFactory());
     _factory->setPeerStorage(_peerStorage);
     _factory->setPeer(_peer);
@@ -155,7 +165,7 @@ void DefaultExtensionMessageFactoryTest::testCreateMessage_UTMetadataRequest()
 void DefaultExtensionMessageFactoryTest::testCreateMessage_UTMetadataData()
 {
   std::string data = getExtensionMessageID("ut_metadata")+
-    "d8:msg_typei1e5:piecei1e10:total_sizei300ee10:0000000000";
+    "d8:msg_typei1e5:piecei1e10:total_sizei300ee0000000000";
   SharedHandle<UTMetadataDataExtensionMessage> m =
     createMessage<UTMetadataDataExtensionMessage>(data);
   CPPUNIT_ASSERT_EQUAL((size_t)1, m->getIndex());
diff --git a/test/HandshakeExtensionMessageTest.cc b/test/HandshakeExtensionMessageTest.cc
index b8ceb4dc..03818f6e 100644
--- a/test/HandshakeExtensionMessageTest.cc
+++ b/test/HandshakeExtensionMessageTest.cc
@@ -9,6 +9,8 @@
 #include "FileEntry.h"
 #include "DownloadContext.h"
 #include "bittorrent_helper.h"
+#include "Option.h"
+#include "RequestGroup.h"
 
 namespace aria2 {
 
@@ -88,9 +90,16 @@ void HandshakeExtensionMessageTest::testToString()
 
 void HandshakeExtensionMessageTest::testDoReceivedAction()
 {
-  SharedHandle<DownloadContext> ctx(new DownloadContext());
+  SharedHandle<DownloadContext> dctx
+    (new DownloadContext(METADATA_PIECE_SIZE, 0));
+  SharedHandle<Option> op(new Option());
+  RequestGroup rg(op);
+  rg.setDownloadContext(dctx);
+  dctx->setOwnerRequestGroup(&rg);
+
   BDE attrs = BDE::dict();
-  ctx->setAttribute(bittorrent::BITTORRENT, attrs);
+  dctx->setAttribute(bittorrent::BITTORRENT, attrs);
+  dctx->markTotalLengthIsUnknown();
 
   SharedHandle<Peer> peer(new Peer("192.168.0.1", 0));
   peer->allocateSessionResource(1024, 1024*1024);
@@ -101,7 +110,7 @@ void HandshakeExtensionMessageTest::testDoReceivedAction()
   msg.setExtension("a2_dht", 2);
   msg.setMetadataSize(1024);
   msg.setPeer(peer);
-  msg.setDownloadContext(ctx);
+  msg.setDownloadContext(dctx);
 
   msg.doReceivedAction();
 
@@ -109,6 +118,8 @@ void HandshakeExtensionMessageTest::testDoReceivedAction()
   CPPUNIT_ASSERT_EQUAL((uint8_t)1, peer->getExtensionMessageID("ut_pex"));
   CPPUNIT_ASSERT_EQUAL((uint8_t)2, peer->getExtensionMessageID("a2_dht"));
   CPPUNIT_ASSERT_EQUAL((int64_t)1024, attrs[bittorrent::METADATA_SIZE].i());
+  CPPUNIT_ASSERT_EQUAL((uint64_t)1024, dctx->getTotalLength());
+  CPPUNIT_ASSERT(dctx->knowsTotalLength());
 }
 
 void HandshakeExtensionMessageTest::testCreate()
diff --git a/test/MockBtMessage.h b/test/MockBtMessage.h
index c0c30353..6feb8d4b 100644
--- a/test/MockBtMessage.h
+++ b/test/MockBtMessage.h
@@ -98,6 +98,15 @@ public:
 
 };
 
+template<typename T>
+class WrapBtMessage:public MockBtMessage {
+public:
+  SharedHandle<T> _m;
+  
+  WrapBtMessage(const SharedHandle<T>& m):_m(m) {}
+};
+
+
 } // namespace aria2
 
 #endif // _D_MOCK_BT_MESSAGE_H_
diff --git a/test/UTMetadataDataExtensionMessageTest.cc b/test/UTMetadataDataExtensionMessageTest.cc
index 463aa3a1..9d8c6ccb 100644
--- a/test/UTMetadataDataExtensionMessageTest.cc
+++ b/test/UTMetadataDataExtensionMessageTest.cc
@@ -5,6 +5,17 @@
 #include <cppunit/extensions/HelperMacros.h>
 
 #include "BtConstants.h"
+#include "PieceStorage.h"
+#include "DownloadContext.h"
+#include "BtRuntime.h"
+#include "DirectDiskAdaptor.h"
+#include "ByteArrayDiskWriter.h"
+#include "BDE.h"
+#include "DownloadContext.h"
+#include "MockPieceStorage.h"
+#include "UTMetadataRequestTracker.h"
+#include "bittorrent_helper.h"
+#include "MessageDigestHelper.h"
 
 namespace aria2 {
 
@@ -41,7 +52,7 @@ void UTMetadataDataExtensionMessageTest::testGetBencodedData()
   msg.setTotalSize(data.size());
   msg.setData(data);
   CPPUNIT_ASSERT_EQUAL
-    (std::string("d8:msg_typei1e5:piecei1e10:total_sizei16384ee16384:")+data,
+    (std::string("d8:msg_typei1e5:piecei1e10:total_sizei16384ee")+data,
      msg.getBencodedData());
 }
 
@@ -55,6 +66,52 @@ void UTMetadataDataExtensionMessageTest::testToString()
 
 void UTMetadataDataExtensionMessageTest::testDoReceivedAction()
 {
+  SharedHandle<DirectDiskAdaptor> diskAdaptor(new DirectDiskAdaptor());
+  SharedHandle<ByteArrayDiskWriter> diskWriter(new ByteArrayDiskWriter());
+  diskAdaptor->setDiskWriter(diskWriter);
+  SharedHandle<MockPieceStorage> pieceStorage(new MockPieceStorage());
+  pieceStorage->setDiskAdaptor(diskAdaptor);
+  SharedHandle<BtRuntime> btRuntime(new BtRuntime());
+  SharedHandle<UTMetadataRequestTracker> tracker
+    (new UTMetadataRequestTracker());
+  SharedHandle<DownloadContext> dctx(new DownloadContext());
+  BDE attrs = BDE::dict();
+
+  std::string piece0 = std::string(METADATA_PIECE_SIZE, '0');
+  std::string piece1 = std::string(METADATA_PIECE_SIZE, '1');
+  std::string metadata = piece0+piece1;
+
+  unsigned char infoHash[INFO_HASH_LENGTH];
+  MessageDigestHelper::digest(infoHash, INFO_HASH_LENGTH,
+			      MessageDigestContext::SHA1,
+			      metadata.data(), metadata.size());
+  attrs[bittorrent::INFO_HASH] = std::string(&infoHash[0], &infoHash[20]);
+
+  dctx->setAttribute(bittorrent::BITTORRENT, attrs);
+
+  UTMetadataDataExtensionMessage m(1);
+  m.setPieceStorage(pieceStorage);
+  m.setBtRuntime(btRuntime);
+  m.setUTMetadataRequestTracker(tracker);
+  m.setDownloadContext(dctx);
+
+  m.setIndex(1);
+  m.setData(piece1);
+  
+  tracker->add(1);
+  m.doReceivedAction();
+  CPPUNIT_ASSERT(!tracker->tracks(1));
+
+  pieceStorage->setDownloadFinished(true);
+  // If piece is not tracked, it is ignored.
+  m.setIndex(0);
+  m.setData(piece0);
+  m.doReceivedAction();
+  CPPUNIT_ASSERT(!btRuntime->isHalt());
+
+  tracker->add(0);
+  m.doReceivedAction();
+  CPPUNIT_ASSERT(btRuntime->isHalt());
 }
 
 } // namespace aria2
diff --git a/test/UTMetadataRequestExtensionMessageTest.cc b/test/UTMetadataRequestExtensionMessageTest.cc
index 6d21d776..e8854b50 100644
--- a/test/UTMetadataRequestExtensionMessageTest.cc
+++ b/test/UTMetadataRequestExtensionMessageTest.cc
@@ -13,6 +13,9 @@
 #include "BtHandshakeMessage.h"
 #include "UTMetadataRejectExtensionMessage.h"
 #include "UTMetadataDataExtensionMessage.h"
+#include "PieceStorage.h"
+#include "BtRuntime.h"
+#include "extension_message_test_helper.h"
 
 namespace aria2 {
 
@@ -27,30 +30,14 @@ class UTMetadataRequestExtensionMessageTest:public CppUnit::TestFixture {
   CPPUNIT_TEST(testDoReceivedAction_data);
   CPPUNIT_TEST_SUITE_END();
 public:
-  class MockExtensionMessage:public MockBtMessage {
-  public:
-    SharedHandle<ExtensionMessage> _m;
-    
-    MockExtensionMessage(const SharedHandle<ExtensionMessage>& m):_m(m) {}
-  };
-
-  class MockBtMessageFactory2:public MockBtMessageFactory {
-  public:
-    virtual SharedHandle<BtMessage>
-    createBtExtendedMessage(const SharedHandle<ExtensionMessage>& extmsg)
-    {
-      return SharedHandle<BtMessage>(new MockExtensionMessage(extmsg));
-    }
-  };
-
   SharedHandle<DownloadContext> _dctx;
-  SharedHandle<MockBtMessageFactory2> _messageFactory;
+  SharedHandle<WrapExtBtMessageFactory> _messageFactory;
   SharedHandle<MockBtMessageDispatcher> _dispatcher;
   SharedHandle<Peer> _peer;
 
   void setUp()
   {
-    _messageFactory.reset(new MockBtMessageFactory2());
+    _messageFactory.reset(new WrapExtBtMessageFactory());
     _dispatcher.reset(new MockBtMessageDispatcher());
     _dctx.reset(new DownloadContext());
     BDE attrs = BDE::dict();
@@ -63,8 +50,8 @@ public:
   template<typename T>
   SharedHandle<T> getFirstDispatchedMessage()
   {
-    SharedHandle<MockExtensionMessage> wrapmsg =
-      dynamic_pointer_cast<MockExtensionMessage>
+    SharedHandle<WrapExtBtMessage> wrapmsg =
+      dynamic_pointer_cast<WrapExtBtMessage>
       (_dispatcher->messageQueue.front());
     
     SharedHandle<T> msg = dynamic_pointer_cast<T>(wrapmsg->_m);
diff --git a/test/UTMetadataRequestFactoryTest.cc b/test/UTMetadataRequestFactoryTest.cc
new file mode 100644
index 00000000..3245e20a
--- /dev/null
+++ b/test/UTMetadataRequestFactoryTest.cc
@@ -0,0 +1,76 @@
+#include "UTMetadataRequestFactory.h"
+
+#include <vector>
+#include <deque>
+
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "MockPieceStorage.h"
+#include "DownloadContext.h"
+#include "Peer.h"
+#include "BtMessage.h"
+#include "extension_message_test_helper.h"
+#include "BtHandshakeMessage.h"
+#include "ExtensionMessage.h"
+#include "UTMetadataRequestTracker.h"
+
+namespace aria2 {
+
+class UTMetadataRequestFactoryTest:public CppUnit::TestFixture {
+
+  CPPUNIT_TEST_SUITE(UTMetadataRequestFactoryTest);
+  CPPUNIT_TEST(testCreate);
+  CPPUNIT_TEST_SUITE_END();
+public:
+  void testCreate();
+
+  class MockPieceStorage2:public MockPieceStorage {
+  public:
+    std::set<size_t> missingIndexes;
+
+    virtual SharedHandle<Piece> getMissingPiece(size_t index)
+    {
+      if(missingIndexes.find(index) != missingIndexes.end()) {
+	return SharedHandle<Piece>(new Piece(index, 0));
+      } else {
+	return SharedHandle<Piece>();
+      }
+    }
+  };
+};
+
+
+CPPUNIT_TEST_SUITE_REGISTRATION(UTMetadataRequestFactoryTest);
+
+void UTMetadataRequestFactoryTest::testCreate()
+{
+  UTMetadataRequestFactory factory;
+  SharedHandle<DownloadContext> dctx
+    (new DownloadContext(METADATA_PIECE_SIZE, METADATA_PIECE_SIZE*2));
+  factory.setDownloadContext(dctx);
+  SharedHandle<MockPieceStorage2> ps(new MockPieceStorage2());
+  ps->missingIndexes.insert(0);
+  ps->missingIndexes.insert(1);
+  SharedHandle<WrapExtBtMessageFactory> messageFactory
+    (new WrapExtBtMessageFactory());
+  factory.setBtMessageFactory(messageFactory);
+  SharedHandle<Peer> peer(new Peer("peer", 6880));
+  peer->allocateSessionResource(0, 0);
+  factory.setPeer(peer);
+  SharedHandle<UTMetadataRequestTracker> tracker
+    (new UTMetadataRequestTracker());
+  factory.setUTMetadataRequestTracker(tracker);
+
+  std::deque<SharedHandle<BtMessage> > msgs;
+
+  factory.create(msgs, 1, ps);
+  CPPUNIT_ASSERT_EQUAL((size_t)1, msgs.size());
+
+  msgs.clear();
+
+  ps->missingIndexes.clear();
+  factory.create(msgs, 1, ps);
+  CPPUNIT_ASSERT_EQUAL((size_t)0, msgs.size());
+}
+
+} // namespace aria2
diff --git a/test/UTMetadataRequestTrackerTest.cc b/test/UTMetadataRequestTrackerTest.cc
new file mode 100644
index 00000000..1170d6dc
--- /dev/null
+++ b/test/UTMetadataRequestTrackerTest.cc
@@ -0,0 +1,72 @@
+#include "UTMetadataRequestTracker.h"
+
+#include <cppunit/extensions/HelperMacros.h>
+
+namespace aria2 {
+
+class UTMetadataRequestTrackerTest:public CppUnit::TestFixture {
+
+  CPPUNIT_TEST_SUITE(UTMetadataRequestTrackerTest);
+  CPPUNIT_TEST(testAdd);
+  CPPUNIT_TEST(testRemove);
+  CPPUNIT_TEST(testGetAllTrackedIndex);
+  CPPUNIT_TEST(testCount);
+  CPPUNIT_TEST(testAvail);
+  CPPUNIT_TEST_SUITE_END();
+public:
+  void testAdd();
+  void testRemove();
+  void testGetAllTrackedIndex();
+  void testCount();
+  void testAvail();
+};
+
+
+CPPUNIT_TEST_SUITE_REGISTRATION(UTMetadataRequestTrackerTest);
+
+void UTMetadataRequestTrackerTest::testAdd()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  CPPUNIT_ASSERT(tr.tracks(1));
+}
+
+void UTMetadataRequestTrackerTest::testRemove()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  tr.remove(1);
+  CPPUNIT_ASSERT(!tr.tracks(1));
+}
+
+void UTMetadataRequestTrackerTest::testGetAllTrackedIndex()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  tr.add(2);
+
+  std::vector<size_t> indexes = tr.getAllTrackedIndex();
+  CPPUNIT_ASSERT_EQUAL((size_t)2, indexes.size());
+  CPPUNIT_ASSERT_EQUAL((size_t)1, indexes[0]);
+  CPPUNIT_ASSERT_EQUAL((size_t)2, indexes[1]);
+}
+
+void UTMetadataRequestTrackerTest::testCount()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  tr.add(2);
+  CPPUNIT_ASSERT_EQUAL((size_t)2, tr.count());
+}
+
+void UTMetadataRequestTrackerTest::testAvail()
+{
+  UTMetadataRequestTracker tr;
+  CPPUNIT_ASSERT_EQUAL((size_t)1, tr.avail());
+  tr.add(1);
+  CPPUNIT_ASSERT_EQUAL((size_t)0, tr.avail());
+  tr.add(2);
+  CPPUNIT_ASSERT_EQUAL((size_t)0, tr.avail());
+}
+
+} // namespace aria2
diff --git a/test/extension_message_test_helper.h b/test/extension_message_test_helper.h
new file mode 100644
index 00000000..54340d20
--- /dev/null
+++ b/test/extension_message_test_helper.h
@@ -0,0 +1,22 @@
+#ifndef _D_EXTENSION_MESSAGE_TEST_HELPER_H_
+#define _D_EXTENSION_MESSAGE_TEST_HELPER_H_
+
+#include "MockBtMessage.h"
+#include "MockBtMessageFactory.h"
+
+namespace aria2 {
+
+typedef WrapBtMessage<ExtensionMessage> WrapExtBtMessage;
+
+class WrapExtBtMessageFactory:public MockBtMessageFactory {
+public:
+  virtual SharedHandle<BtMessage>
+  createBtExtendedMessage(const SharedHandle<ExtensionMessage>& extmsg)
+  {
+    return SharedHandle<BtMessage>(new WrapExtBtMessage(extmsg));
+  }
+};
+
+} // namespace aria2
+
+#endif // _D_EXTENSION_MESSAGE_TEST_HELPER_H_