Rewritten ExtensionMessageRegistry

1 files changed, 21 insertions, 14 deletions

diff --git a/test/HandshakeExtensionMessageTest.cc b/test/HandshakeExtensionMessageTest.cc
index 064a006e..eea01d21 100644
--- a/test/HandshakeExtensionMessageTest.cc
+++ b/test/HandshakeExtensionMessageTest.cc
@@ -60,12 +60,12 @@ void HandshakeExtensionMessageTest::testGetBencodedData()
   HandshakeExtensionMessage msg;
   msg.setClientVersion("aria2");
   msg.setTCPPort(6889);
-  msg.setExtension("ut_pex", 1);
-  msg.setExtension("a2_dht", 2);
+  msg.setExtension(ExtensionMessageRegistry::UT_PEX, 1);
+  msg.setExtension(ExtensionMessageRegistry::UT_METADATA, 2);
   msg.setMetadataSize(1024);
   CPPUNIT_ASSERT_EQUAL
     (std::string("d"
-                 "1:md6:a2_dhti2e6:ut_pexi1ee"
+                 "1:md11:ut_metadatai2e6:ut_pexi1ee"
                  "13:metadata_sizei1024e"
                  "1:pi6889e"
                  "1:v5:aria2"
@@ -81,12 +81,12 @@ void HandshakeExtensionMessageTest::testToString()
   HandshakeExtensionMessage msg;
   msg.setClientVersion("aria2");
   msg.setTCPPort(6889);
-  msg.setExtension("ut_pex", 1);
-  msg.setExtension("a2_dht", 2);
+  msg.setExtension(ExtensionMessageRegistry::UT_PEX, 1);
+  msg.setExtension(ExtensionMessageRegistry::UT_METADATA, 2);
   msg.setMetadataSize(1024);
   CPPUNIT_ASSERT_EQUAL
     (std::string("handshake client=aria2, tcpPort=6889, metadataSize=1024,"
-                 " a2_dht=2, ut_pex=1"), msg.toString());
+                 " ut_metadata=2, ut_pex=1"), msg.toString());
 }
 
 void HandshakeExtensionMessageTest::testDoReceivedAction()
@@ -106,9 +106,8 @@ void HandshakeExtensionMessageTest::testDoReceivedAction()
   HandshakeExtensionMessage msg;
   msg.setClientVersion("aria2");
   msg.setTCPPort(6889);
-  msg.setExtension("ut_pex", 1);
-  msg.setExtension("a2_dht", 2);
-  msg.setExtension("ut_metadata", 3);
+  msg.setExtension(ExtensionMessageRegistry::UT_PEX, 1);
+  msg.setExtension(ExtensionMessageRegistry::UT_METADATA, 3);
   msg.setMetadataSize(1024);
   msg.setPeer(peer);
   msg.setDownloadContext(dctx);
@@ -116,8 +115,12 @@ void HandshakeExtensionMessageTest::testDoReceivedAction()
   msg.doReceivedAction();
 
   CPPUNIT_ASSERT_EQUAL((uint16_t)6889, peer->getPort());
-  CPPUNIT_ASSERT_EQUAL((uint8_t)1, peer->getExtensionMessageID("ut_pex"));
-  CPPUNIT_ASSERT_EQUAL((uint8_t)2, peer->getExtensionMessageID("a2_dht"));
+  CPPUNIT_ASSERT_EQUAL((uint8_t)1,
+                       peer->getExtensionMessageID
+                       (ExtensionMessageRegistry::UT_PEX));
+  CPPUNIT_ASSERT_EQUAL((uint8_t)3,
+                       peer->getExtensionMessageID
+                       (ExtensionMessageRegistry::UT_METADATA));
   CPPUNIT_ASSERT(peer->isSeeder());
   CPPUNIT_ASSERT_EQUAL((size_t)1024, attrs->metadataSize);
   CPPUNIT_ASSERT_EQUAL((int64_t)1024, dctx->getTotalLength());
@@ -134,13 +137,15 @@ void HandshakeExtensionMessageTest::testDoReceivedAction()
 void HandshakeExtensionMessageTest::testCreate()
 {
   std::string in = 
-    "0d1:pi6881e1:v5:aria21:md6:ut_pexi1ee13:metadata_sizei1024ee";
+    "0d1:pi6881e1:v5:aria21:md5:a2dhti2e6:ut_pexi1ee13:metadata_sizei1024ee";
   SharedHandle<HandshakeExtensionMessage> m =
     HandshakeExtensionMessage::create(reinterpret_cast<const unsigned char*>(in.c_str()),
                                       in.size());
   CPPUNIT_ASSERT_EQUAL(std::string("aria2"), m->getClientVersion());
   CPPUNIT_ASSERT_EQUAL((uint16_t)6881, m->getTCPPort());
-  CPPUNIT_ASSERT_EQUAL((uint8_t)1, m->getExtensionMessageID("ut_pex"));
+  CPPUNIT_ASSERT_EQUAL((uint8_t)1,
+                       m->getExtensionMessageID
+                       (ExtensionMessageRegistry::UT_PEX));
   CPPUNIT_ASSERT_EQUAL((size_t)1024, m->getMetadataSize());
   try {
     // bad payload format
@@ -182,7 +187,9 @@ void HandshakeExtensionMessageTest::testCreate_stringnum()
   // port number in string is not allowed
   CPPUNIT_ASSERT_EQUAL((uint16_t)0, m->getTCPPort());
   // extension ID in string is not allowed
-  CPPUNIT_ASSERT_EQUAL((uint8_t)0, m->getExtensionMessageID("ut_pex"));
+  CPPUNIT_ASSERT_EQUAL((uint8_t)0,
+                       m->getExtensionMessageID
+                       (ExtensionMessageRegistry::UT_PEX));
 }
 
 } // namespace aria2