did basic arbitrage service tests

arbitrage_engine/src/arbitrage_service.rs

@@ -1,7 +1,8 @@
-use super::decorated_edge::DecoratedEdge;
-use super::engine::{get_all_negative_cycles_0, get_negative_cycle_quick};
-use super::path::{DecoratedPath, Path};
-
+use crate::{
+    decorated_edge::DecoratedEdge,
+    engine::{get_all_negative_cycles_0, get_negative_cycle_quick},
+    path::{DecoratedPath, Path}
+};
 use blake3::Hash;
 use petgraph::{
     graph::{EdgeReference, Graph},
@@ -10,8 +11,7 @@ use petgraph::{
 };
 use std::collections::HashMap;
 
-trait IArbitrageService {
-    fn new() -> Self;
+pub trait IArbitrageService {
     fn upsert_path(&mut self, n0: &str, n1: &str, edge: DecoratedEdge) -> bool;
     /// Returns all arbitrages found
     fn scan_arbitrages(&self) -> Vec<DecoratedPath>;
@@ -22,7 +22,7 @@ trait IArbitrageService {
 }
 
 /// Point of contact interacting with the arbitrage functionality
-struct ArbitrageService {
+pub struct ArbitrageService {
     /// Core directed graph data structure on which we perform the algorithm
     /// &str for nodeId, f64 for edge weight
     graph: Graph<String, f64>,
@@ -37,8 +37,8 @@ struct ArbitrageService {
     decorated_edges: HashMap<EdgeIndex, DecoratedEdge>,
 }
 
-impl IArbitrageService for ArbitrageService {
-    fn new() -> Self {
+impl ArbitrageService {
+    pub fn new() -> Self {
         ArbitrageService {
             graph: Graph::new(),
             node_indexes: HashMap::new(),
@@ -48,6 +48,22 @@ impl IArbitrageService for ArbitrageService {
         }
     }
 
+    pub fn node_count(&self) -> usize {
+        self.graph.node_count()
+    }
+
+    pub fn edge_count(&self) -> usize {
+        self.graph.node_count()
+    }
+}
+
+impl Default for ArbitrageService {
+    fn default() -> Self {
+        ArbitrageService::new()
+    }
+}
+
+impl IArbitrageService for ArbitrageService {
     /// Add edge from n0 -> n1
     /// We will not add the reverse edge here - up to the client
     /// Return true if graph updated, false if not

arbitrage_engine/src/arbitrage_service_test.rs

@@ -0,0 +1,226 @@
+#[cfg(test)]
+mod tests {
+    use crate::{
+        decorated_edge::DecoratedEdge,
+        arbitrage_service::{IArbitrageService, ArbitrageService}
+    };
+
+    #[test]
+    fn test_upsert_path_single_path_success() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+        let edge = DecoratedEdge {
+            weight: 1.0,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let is_updated = service.upsert_path("a", "b", edge);
+        assert!(is_updated);
+    }
+
+    #[test]
+    fn test_upsert_path_duplicate_path_should_return_false() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+        let edge = DecoratedEdge {
+            weight: 1.0,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let duplicated_edge = DecoratedEdge {
+            weight: 1.0,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let _ = service.upsert_path("a", "b", edge);
+        let is_updated = service.upsert_path("a", "b", duplicated_edge);
+        assert!(!is_updated);
+    }
+
+    #[test]
+    fn test_upsert_path_two_paths_between_same_nodes_success() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+        let edge_0 = DecoratedEdge {
+            weight: 1.0,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let edge_1 = DecoratedEdge {
+            weight: 1.0,
+            protocol_type: 2,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let _ = service.upsert_path("a", "b", edge_0);
+        let is_updated = service.upsert_path("a", "b", edge_1);
+        assert!(is_updated);
+    }
+
+    #[test]
+    fn test_upsert_path_third_edge_between_existing_edges_success() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+        let edge_0 = DecoratedEdge {
+            weight: 0.95,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let edge_1 = DecoratedEdge {
+            weight: 1.05,
+            protocol_type: 2,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let new_edge = DecoratedEdge {
+            weight: 1.0,
+            protocol_type: 3,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let _ = service.upsert_path("a", "b", edge_0);
+        let _ = service.upsert_path("a", "b", edge_1);
+        let is_updated = service.upsert_path("a", "b", new_edge);
+        assert!(!is_updated);
+    }
+
+    #[test]
+    fn test_upsert_path_third_edge_above_existing_edges_success() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+        let edge_0 = DecoratedEdge {
+            weight: 0.95,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let edge_1 = DecoratedEdge {
+            weight: 1.05,
+            protocol_type: 2,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let new_edge = DecoratedEdge {
+            weight: 1.10,
+            protocol_type: 3,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let _ = service.upsert_path("a", "b", edge_0);
+        let _ = service.upsert_path("a", "b", edge_1);
+        let is_updated = service.upsert_path("a", "b", new_edge);
+        assert!(is_updated);
+    }
+
+    #[test]
+    fn test_upsert_path_third_edge_below_existing_edges_success() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+        let edge_0 = DecoratedEdge {
+            weight: 0.95,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let edge_1 = DecoratedEdge {
+            weight: 1.05,
+            protocol_type: 2,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let new_edge = DecoratedEdge {
+            weight: 0.90,
+            protocol_type: 3,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+        let _ = service.upsert_path("a", "b", edge_0);
+        let _ = service.upsert_path("a", "b", edge_1);
+        let is_updated = service.upsert_path("a", "b", new_edge);
+        assert!(is_updated);
+    }
+
+    // Use same test conditions as engine_test.get_negative_cycle_for_source_quick_test_2
+    #[test]
+    fn test_scan_arbitrages_quick_success() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+
+        let new_decorated_edge = |weight: f64| DecoratedEdge {
+            weight,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+
+        service.upsert_path("nodes[4]", "nodes[5]", new_decorated_edge(0.35));
+        service.upsert_path("nodes[5]", "nodes[4]", new_decorated_edge(-0.66));
+        service.upsert_path("nodes[4]", "nodes[7]", new_decorated_edge(0.37));
+        service.upsert_path("nodes[5]", "nodes[7]", new_decorated_edge(0.28));
+        service.upsert_path("nodes[7]", "nodes[5]", new_decorated_edge(0.28));
+        service.upsert_path("nodes[5]", "nodes[1]", new_decorated_edge(0.32));
+        service.upsert_path("nodes[0]", "nodes[4]", new_decorated_edge(0.38));
+        service.upsert_path("nodes[0]", "nodes[2]", new_decorated_edge(0.26));
+        service.upsert_path("nodes[7]", "nodes[3]", new_decorated_edge(0.39));
+        service.upsert_path("nodes[1]", "nodes[3]", new_decorated_edge(0.29));
+        service.upsert_path("nodes[2]", "nodes[7]", new_decorated_edge(0.34));
+        service.upsert_path("nodes[6]", "nodes[2]", new_decorated_edge(0.40));
+        service.upsert_path("nodes[3]", "nodes[6]", new_decorated_edge(0.52));
+        service.upsert_path("nodes[6]", "nodes[0]", new_decorated_edge(0.58));
+        service.upsert_path("nodes[6]", "nodes[4]", new_decorated_edge(0.93));
+
+        assert_eq!(service.node_count(), 8);
+        let paths = service.scan_arbitrages_quick();
+        assert_eq!(paths.len(), 1);
+        assert_eq!(paths[0].nodes.len(), 3);
+        assert_eq!(paths[0].edges.len(), 2);
+    }
+
+    #[test]
+    fn test_scan_arbitrages_quick() {
+        let mut service: ArbitrageService = ArbitrageService::new();
+
+        let new_decorated_edge = |weight: f64| DecoratedEdge {
+            weight,
+            protocol_type: 1,
+            node_type: 1,
+            pool_address: "".to_string(),
+            data: "".to_string()
+        };
+
+        service.upsert_path("nodes[4]", "nodes[5]", new_decorated_edge(0.35));
+        service.upsert_path("nodes[5]", "nodes[4]", new_decorated_edge(-0.66));
+        service.upsert_path("nodes[4]", "nodes[7]", new_decorated_edge(0.37));
+        service.upsert_path("nodes[5]", "nodes[7]", new_decorated_edge(0.28));
+        service.upsert_path("nodes[7]", "nodes[5]", new_decorated_edge(0.28));
+        service.upsert_path("nodes[5]", "nodes[1]", new_decorated_edge(0.32));
+        service.upsert_path("nodes[0]", "nodes[4]", new_decorated_edge(0.38));
+        service.upsert_path("nodes[0]", "nodes[2]", new_decorated_edge(0.26));
+        service.upsert_path("nodes[7]", "nodes[3]", new_decorated_edge(0.39));
+        service.upsert_path("nodes[1]", "nodes[3]", new_decorated_edge(0.29));
+        service.upsert_path("nodes[2]", "nodes[7]", new_decorated_edge(0.34));
+        service.upsert_path("nodes[6]", "nodes[2]", new_decorated_edge(0.40));
+        service.upsert_path("nodes[3]", "nodes[6]", new_decorated_edge(0.52));
+        service.upsert_path("nodes[6]", "nodes[0]", new_decorated_edge(0.58));
+        service.upsert_path("nodes[6]", "nodes[4]", new_decorated_edge(0.93));
+
+        assert_eq!(service.node_count(), 8);
+        let paths = service.scan_arbitrages();
+        assert_eq!(paths.len(), 1);
+        assert_eq!(paths[0].nodes.len(), 3);
+        assert_eq!(paths[0].edges.len(), 2);
+    }
+}

arbitrage_engine/src/decorated_edge.rs

@@ -8,18 +8,19 @@ use blake3::{Hash, Hasher};
 pub struct DecoratedEdge {
     pub weight: f64,
     /// int enum for protocol
-    protocol_type: usize,
+    pub protocol_type: usize,
     /// int enum for blockchain node type (e.g. EVM, Solana)
-    node_type: usize,
+    pub node_type: usize,
     /// Hex string address for pool
-    pool_address: String,
+    pub pool_address: String,
     /// Miscellaneous data
-    data: String,
+    pub data: String,
 }
 
 impl DecoratedEdge {
     pub fn get_unique_id(edge: &DecoratedEdge) -> Hash {
         let mut hasher = Hasher::new();
+        hasher.update(&edge.weight.to_ne_bytes());
         hasher.update(&edge.protocol_type.to_ne_bytes());
         hasher.update(&edge.node_type.to_ne_bytes());
         hasher.update(edge.pool_address.as_bytes());