Verify that we only download missing piece #32 (#58)

Zitrax · Nov 12, 2022 · 4d9f355 · 4d9f355
1 parent ce2a7e1
commit 4d9f355
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 5 deletions.
diff --git a/src/file_writer.hpp b/src/file_writer.hpp
@@ -110,7 +110,7 @@ class FileWriterThread {
     m_file_writer.set_callback(std::move(cb));
     using std::placeholders::_1;
     using std::placeholders::_2;
-    torrent.set_piece_callback(
+    torrent.add_piece_callback(
         [&](Torrent* t, const std::shared_ptr<Piece>& piece) {
           m_file_writer.add(t, piece);
         });

diff --git a/src/torrent.cpp b/src/torrent.cpp
@@ -681,7 +681,7 @@ void Torrent::last_piece_written() {
 
 void Torrent::piece_done(std::shared_ptr<Piece>& piece) {
   m_client_pieces[piece->id()] = true;
-  m_piece_callback(this, piece);
+  ranges::for_each(m_piece_callbacks, [&](const auto& cb) { cb(this, piece); });
 }
 
 ostream& operator<<(ostream& os, const zit::FileInfo& file_info) {

diff --git a/src/torrent.hpp b/src/torrent.hpp
@@ -174,8 +174,8 @@ class Torrent {
   /**
    * Callback that will be called whenever a piece has finished downloading.
    */
-  void set_piece_callback(PieceCallback piece_callback) {
-    m_piece_callback = std::move(piece_callback);
+  void add_piece_callback(PieceCallback piece_callback) {
+    m_piece_callbacks.emplace_back(std::move(piece_callback));
   }
 
   /**
@@ -338,7 +338,7 @@ class Torrent {
   std::shared_ptr<spdlog::logger> m_logger{};
   std::filesystem::path m_tmpfile{};
   std::filesystem::path m_data_dir{};
-  PieceCallback m_piece_callback{};
+  std::vector<PieceCallback> m_piece_callbacks{};
   DisconnectCallback m_disconnect_callback{};
   std::string m_peer_id{};
   ListeningPort m_listening_port;

diff --git a/tests/test_integrate.cpp b/tests/test_integrate.cpp
@@ -285,13 +285,22 @@ TEST_P(IntegrateF, DISABLED_download_part) {
   TestConfig test_config;
   zit::Torrent torrent(torrent_file, download_dir, test_config);
   ASSERT_FALSE(torrent.done());
+
+  unsigned pieces_downloaded = 0;
+  torrent.add_piece_callback(
+      [&pieces_downloaded](auto, auto) { pieces_downloaded++; });
+
   auto target = download(data_dir, torrent_file, torrent, 1);
 
   // Transfer done - Verify content
   auto source = data_dir / "1MiB.dat";
   auto source_sha1 = zit::Sha1::calculateFile(source).hex();
   auto target_sha1 = zit::Sha1::calculateFile(target).hex();
   EXPECT_EQ(source_sha1, target_sha1);
+
+  // Since we only changed one piece we should only have to get one piece again
+  EXPECT_GT(torrent.pieces().size(), 1);
+  EXPECT_EQ(pieces_downloaded, 1);
 }
 
 #ifdef INTEGRATION_TESTS
@@ -320,6 +329,11 @@ TEST_F(Integrate, DISABLED_download_multi_part) {
   TestConfig test_config;
   zit::Torrent torrent(torrent_file, download_dir, test_config);
   ASSERT_FALSE(torrent.done());
+
+  unsigned pieces_downloaded = 0;
+  torrent.add_piece_callback(
+      [&pieces_downloaded](auto, auto) { pieces_downloaded++; });
+
   auto target = download(data_dir, torrent_file, torrent, 1);
 
   // Transfer done - Verify content
@@ -334,6 +348,10 @@ TEST_F(Integrate, DISABLED_download_multi_part) {
     fs::remove(dst);
   }
   fs::remove(name);
+
+  // Since we only changed one piece we should only have to get one piece again
+  EXPECT_GT(torrent.pieces().size(), 1);
+  EXPECT_EQ(pieces_downloaded, 1);
 }
 
 #ifdef INTEGRATION_TESTS