From 4d9f3556910dfe2e62bc906a7e9479900d640699 Mon Sep 17 00:00:00 2001
From: Daniel Bengtsson <daniel@bengtssons.info>
Date: Sat, 12 Nov 2022 12:11:08 +0100
Subject: [PATCH] Verify that we only download missing piece #32 (#58)

---
 src/file_writer.hpp      |  2 +-
 src/torrent.cpp          |  2 +-
 src/torrent.hpp          |  6 +++---
 tests/test_integrate.cpp | 18 ++++++++++++++++++
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/file_writer.hpp b/src/file_writer.hpp
index 93a039e..98b3cbf 100644
--- 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
index c2f022d..9b7dff4 100644
--- 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
index 258e7d0..941e1a3 100644
--- 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
index 7a35f0e..dc173d0 100644
--- a/tests/test_integrate.cpp
+++ b/tests/test_integrate.cpp
@@ -285,6 +285,11 @@ 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
@@ -292,6 +297,10 @@ TEST_P(IntegrateF, DISABLED_download_part) {
   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