Merge pull request #2110 from munjesi/docs_fixes
[rust-lightning] / lightning / src / util / test_utils.rs
index 2d68e74c2e594978ca645bc5ef1dd7f34466f860..ddc09f0b37d8d3a462755cc78e5a32a5bbb74ac7 100644 (file)
@@ -22,10 +22,10 @@ use crate::ln::features::{ChannelFeatures, InitFeatures, NodeFeatures};
 use crate::ln::{msgs, wire};
 use crate::ln::msgs::LightningError;
 use crate::ln::script::ShutdownScript;
-use crate::routing::gossip::{NetworkGraph, NodeId};
+use crate::routing::gossip::{EffectiveCapacity, NetworkGraph, NodeId};
 use crate::routing::utxo::{UtxoLookup, UtxoLookupError, UtxoResult};
 use crate::routing::router::{find_route, InFlightHtlcs, Route, RouteHop, RouteParameters, Router, ScorerAccountingForInFlightHtlcs};
-use crate::routing::scoring::FixedPenaltyScorer;
+use crate::routing::scoring::{ChannelUsage, Score};
 use crate::util::config::UserConfig;
 use crate::util::enforcing_trait_impls::{EnforcingSigner, EnforcementState};
 use crate::util::events;
@@ -48,6 +48,7 @@ use regex;
 
 use crate::io;
 use crate::prelude::*;
+use core::cell::RefCell;
 use core::time::Duration;
 use crate::sync::{Mutex, Arc};
 use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
@@ -79,11 +80,12 @@ impl chaininterface::FeeEstimator for TestFeeEstimator {
 pub struct TestRouter<'a> {
        pub network_graph: Arc<NetworkGraph<&'a TestLogger>>,
        pub next_routes: Mutex<VecDeque<(RouteParameters, Result<Route, LightningError>)>>,
+       pub scorer: &'a Mutex<TestScorer>,
 }
 
 impl<'a> TestRouter<'a> {
-       pub fn new(network_graph: Arc<NetworkGraph<&'a TestLogger>>) -> Self {
-               Self { network_graph, next_routes: Mutex::new(VecDeque::new()), }
+       pub fn new(network_graph: Arc<NetworkGraph<&'a TestLogger>>, scorer: &'a Mutex<TestScorer>) -> Self {
+               Self { network_graph, next_routes: Mutex::new(VecDeque::new()), scorer }
        }
 
        pub fn expect_find_route(&self, query: RouteParameters, result: Result<Route, LightningError>) {
@@ -99,19 +101,40 @@ impl<'a> Router for TestRouter<'a> {
        ) -> Result<Route, msgs::LightningError> {
                if let Some((find_route_query, find_route_res)) = self.next_routes.lock().unwrap().pop_front() {
                        assert_eq!(find_route_query, *params);
+                       if let Ok(ref route) = find_route_res {
+                               let locked_scorer = self.scorer.lock().unwrap();
+                               let scorer = ScorerAccountingForInFlightHtlcs::new(locked_scorer, inflight_htlcs);
+                               for path in &route.paths {
+                                       let mut aggregate_msat = 0u64;
+                                       for (idx, hop) in path.iter().rev().enumerate() {
+                                               aggregate_msat += hop.fee_msat;
+                                               let usage = ChannelUsage {
+                                                       amount_msat: aggregate_msat,
+                                                       inflight_htlc_msat: 0,
+                                                       effective_capacity: EffectiveCapacity::Unknown,
+                                               };
+
+                                               // Since the path is reversed, the last element in our iteration is the first
+                                               // hop.
+                                               if idx == path.len() - 1 {
+                                                       scorer.channel_penalty_msat(hop.short_channel_id, &NodeId::from_pubkey(payer), &NodeId::from_pubkey(&hop.pubkey), usage);
+                                               } else {
+                                                       let curr_hop_path_idx = path.len() - 1 - idx;
+                                                       scorer.channel_penalty_msat(hop.short_channel_id, &NodeId::from_pubkey(&path[curr_hop_path_idx - 1].pubkey), &NodeId::from_pubkey(&hop.pubkey), usage);
+                                               }
+                                       }
+                               }
+                       }
                        return find_route_res;
                }
                let logger = TestLogger::new();
+               let scorer = self.scorer.lock().unwrap();
                find_route(
                        payer, params, &self.network_graph, first_hops, &logger,
-                       &ScorerAccountingForInFlightHtlcs::new(TestScorer::with_penalty(0), &inflight_htlcs),
+                       &ScorerAccountingForInFlightHtlcs::new(scorer, &inflight_htlcs),
                        &[42; 32]
                )
        }
-       fn notify_payment_path_failed(&self, _path: &[&RouteHop], _short_channel_id: u64) {}
-       fn notify_payment_path_successful(&self, _path: &[&RouteHop]) {}
-       fn notify_payment_probe_successful(&self, _path: &[&RouteHop]) {}
-       fn notify_payment_probe_failed(&self, _path: &[&RouteHop], _short_channel_id: u64) {}
 }
 
 impl<'a> Drop for TestRouter<'a> {
@@ -308,6 +331,7 @@ impl chaininterface::BroadcasterInterface for TestBroadcaster {
 pub struct TestChannelMessageHandler {
        pub pending_events: Mutex<Vec<events::MessageSendEvent>>,
        expected_recv_msgs: Mutex<Option<Vec<wire::Message<()>>>>,
+       connected_peers: Mutex<HashSet<PublicKey>>,
 }
 
 impl TestChannelMessageHandler {
@@ -315,6 +339,7 @@ impl TestChannelMessageHandler {
                TestChannelMessageHandler {
                        pending_events: Mutex::new(Vec::new()),
                        expected_recv_msgs: Mutex::new(None),
+                       connected_peers: Mutex::new(HashSet::new()),
                }
        }
 
@@ -399,8 +424,11 @@ impl msgs::ChannelMessageHandler for TestChannelMessageHandler {
        fn handle_channel_reestablish(&self, _their_node_id: &PublicKey, msg: &msgs::ChannelReestablish) {
                self.received_msg(wire::Message::ChannelReestablish(msg.clone()));
        }
-       fn peer_disconnected(&self, _their_node_id: &PublicKey, _no_connection_possible: bool) {}
-       fn peer_connected(&self, _their_node_id: &PublicKey, _msg: &msgs::Init) -> Result<(), ()> {
+       fn peer_disconnected(&self, their_node_id: &PublicKey) {
+               assert!(self.connected_peers.lock().unwrap().remove(their_node_id));
+       }
+       fn peer_connected(&self, their_node_id: &PublicKey, _msg: &msgs::Init, _inbound: bool) -> Result<(), ()> {
+               assert!(self.connected_peers.lock().unwrap().insert(their_node_id.clone()));
                // Don't bother with `received_msg` for Init as its auto-generated and we don't want to
                // bother re-generating the expected Init message in all tests.
                Ok(())
@@ -516,7 +544,7 @@ impl msgs::RoutingMessageHandler for TestRoutingMessageHandler {
                None
        }
 
-       fn peer_connected(&self, their_node_id: &PublicKey, init_msg: &msgs::Init) -> Result<(), ()> {
+       fn peer_connected(&self, their_node_id: &PublicKey, init_msg: &msgs::Init, _inbound: bool) -> Result<(), ()> {
                if !init_msg.features.supports_gossip_queries() {
                        return Ok(());
                }
@@ -614,10 +642,10 @@ impl TestLogger {
        /// 1. belongs to the specified module and
        /// 2. contains `line` in it.
        /// And asserts if the number of occurrences is the same with the given `count`
-       pub fn assert_log_contains(&self, module: String, line: String, count: usize) {
+       pub fn assert_log_contains(&self, module: &str, line: &str, count: usize) {
                let log_entries = self.lines.lock().unwrap();
                let l: usize = log_entries.iter().filter(|&(&(ref m, ref l), _c)| {
-                       m == &module && l.contains(line.as_str())
+                       m == module && l.contains(line)
                }).map(|(_, c) | { c }).sum();
                assert_eq!(l, count)
        }
@@ -626,10 +654,10 @@ impl TestLogger {
        /// 1. belong to the specified module and
        /// 2. match the given regex pattern.
        /// Assert that the number of occurrences equals the given `count`
-       pub fn assert_log_regex(&self, module: String, pattern: regex::Regex, count: usize) {
+       pub fn assert_log_regex(&self, module: &str, pattern: regex::Regex, count: usize) {
                let log_entries = self.lines.lock().unwrap();
                let l: usize = log_entries.iter().filter(|&(&(ref m, ref l), _c)| {
-                       m == &module && pattern.is_match(&l)
+                       m == module && pattern.is_match(&l)
                }).map(|(_, c) | { c }).sum();
                assert_eq!(l, count)
        }
@@ -890,5 +918,65 @@ impl Drop for TestChainSource {
        }
 }
 
-/// A scorer useful in testing, when the passage of time isn't a concern.
-pub type TestScorer = FixedPenaltyScorer;
+pub struct TestScorer {
+       /// Stores a tuple of (scid, ChannelUsage)
+       scorer_expectations: RefCell<Option<VecDeque<(u64, ChannelUsage)>>>,
+}
+
+impl TestScorer {
+       pub fn new() -> Self {
+               Self {
+                       scorer_expectations: RefCell::new(None),
+               }
+       }
+
+       pub fn expect_usage(&self, scid: u64, expectation: ChannelUsage) {
+               self.scorer_expectations.borrow_mut().get_or_insert_with(|| VecDeque::new()).push_back((scid, expectation));
+       }
+}
+
+#[cfg(c_bindings)]
+impl crate::util::ser::Writeable for TestScorer {
+       fn write<W: crate::util::ser::Writer>(&self, _: &mut W) -> Result<(), crate::io::Error> { unreachable!(); }
+}
+
+impl Score for TestScorer {
+       fn channel_penalty_msat(
+               &self, short_channel_id: u64, _source: &NodeId, _target: &NodeId, usage: ChannelUsage
+       ) -> u64 {
+               if let Some(scorer_expectations) = self.scorer_expectations.borrow_mut().as_mut() {
+                       match scorer_expectations.pop_front() {
+                               Some((scid, expectation)) => {
+                                       assert_eq!(expectation, usage);
+                                       assert_eq!(scid, short_channel_id);
+                               },
+                               None => {},
+                       }
+               }
+               0
+       }
+
+       fn payment_path_failed(&mut self, _actual_path: &[&RouteHop], _actual_short_channel_id: u64) {}
+
+       fn payment_path_successful(&mut self, _actual_path: &[&RouteHop]) {}
+
+       fn probe_failed(&mut self, _actual_path: &[&RouteHop], _: u64) {}
+
+       fn probe_successful(&mut self, _actual_path: &[&RouteHop]) {}
+}
+
+impl Drop for TestScorer {
+       fn drop(&mut self) {
+               #[cfg(feature = "std")] {
+                       if std::thread::panicking() {
+                               return;
+                       }
+               }
+
+               if let Some(scorer_expectations) = self.scorer_expectations.borrow().as_ref() {
+                       if !scorer_expectations.is_empty() {
+                               panic!("Unsatisfied scorer expectations: {:?}", scorer_expectations)
+                       }
+               }
+       }
+}