deduplicate first in new_with_addrs test function (#6724)

ZcashFoundation · May 19, 2023 · a972144 · a972144
1 parent 943d213
commit a972144
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/zebra-network/src/address_book.rs b/zebra-network/src/address_book.rs
@@ -180,6 +180,9 @@ impl AddressBook {
         let instant_now = Instant::now();
         let chrono_now = Utc::now();
 
+        // The maximum number of addresses should be always greater than 0
+        assert!(addr_limit > 0);
+
         let mut new_book = AddressBook::new(local_listener, network, span);
         new_book.addr_limit = addr_limit;
 
@@ -190,12 +193,15 @@ impl AddressBook {
                 meta_addr
             })
             .filter(|meta_addr| meta_addr.address_is_valid_for_outbound(network))
-            .take(addr_limit)
             .map(|meta_addr| (meta_addr.addr, meta_addr));
 
         for (socket_addr, meta_addr) in addrs {
             // overwrite any duplicate addresses
             new_book.by_addr.insert(socket_addr, meta_addr);
+            // exit as soon as we get enough addresses
+            if new_book.by_addr.len() >= addr_limit {
+                break;
+            }
         }
 
         new_book.update_metrics(instant_now, chrono_now);

diff --git a/zebra-network/src/address_book/tests/prop.rs b/zebra-network/src/address_book/tests/prop.rs
@@ -77,7 +77,7 @@ proptest! {
             MetaAddrChange::addr_changes_strategy(MAX_ADDR_CHANGE),
             2..MAX_ADDR_CHANGE
         ),
-        addr_limit in 0..=MAX_ADDR_CHANGE,
+        addr_limit in 1..=MAX_ADDR_CHANGE,
         pre_fill in any::<bool>(),
     ) {
         let _init_guard = zebra_test::init();