diff --git a/src/flow/flow.rs b/src/flow/flow.rs
index d66bdf7..1552909 100644
--- a/src/flow/flow.rs
+++ b/src/flow/flow.rs
@@ -9,7 +9,7 @@ pub fn compute_flow(
     source: &Address,
     sink: &Address,
     edges: &HashMap<Address, Vec<Edge>>,
-) -> String {
+) -> (U256, Vec<Edge>) {
     let mut adjacencies = Adjacencies::new(edges);
     let mut used_edges: HashMap<Node, HashMap<Node, U256>> = HashMap::new();
 
@@ -52,7 +52,7 @@ pub fn compute_flow(
         extract_transfers(source, sink, &flow, used_edges)
     };
     println!("Num transfers: {}", transfers.len());
-    flow.to_string()
+    (flow, transfers)
 }
 
 fn augmenting_path(
@@ -165,3 +165,69 @@ fn next_full_capacity_edge(
     }
     panic!();
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn addresses() -> (Address, Address, Address, Address, Address, Address) {
+        (
+            Address::from("0x11C7e86fF693e9032A0F41711b5581a04b26Be2E"),
+            Address::from("0x22cEDde51198D1773590311E2A340DC06B24cB37"),
+            Address::from("0x33cEDde51198D1773590311E2A340DC06B24cB37"),
+            Address::from("0x447EDde51198D1773590311E2A340DC06B24cB37"),
+            Address::from("0x55c16ce62d26fd51582a646e2e30a3267b1e6d7e"),
+            Address::from("0x66c16ce62d26fd51582a646e2e30a3267b1e6d7e"),
+        )
+    }
+    fn build_edges(input: Vec<Edge>) -> HashMap<Address, Vec<Edge>> {
+        let mut output: HashMap<Address, Vec<Edge>> = HashMap::new();
+        for e in input {
+            output.entry(e.from).or_default().push(e);
+        }
+        output
+    }
+
+    #[test]
+    fn direct() {
+        let (a, b, t, ..) = addresses();
+        let edges = build_edges(vec![
+            Edge{from: a, to: b, token: t, capacity: U256::from(10)}
+        ]);
+        let flow = compute_flow(&a, &b, &edges);
+        assert_eq!(flow, (U256::from(10), edges[&a].clone()));
+    }
+
+    #[test]
+    fn one_hop() {
+        let (a, b, c, t1, t2, ..) = addresses();
+        let edges = build_edges(vec![
+            Edge{from: a, to: b, token: t1, capacity: U256::from(10)},
+            Edge{from: b, to: c, token: t2, capacity: U256::from(8)},
+        ]);
+        let flow = compute_flow(&a, &c, &edges);
+        assert_eq!(flow, (U256::from(8), vec![
+            Edge{from: a, to: b, token: t1, capacity: U256::from(8)},
+            Edge{from: b, to: c, token: t2, capacity: U256::from(8)},
+        ]));
+    }
+
+    #[test]
+    fn diamond() {
+        let (a, b, c, d, t1, t2) = addresses();
+        let edges = build_edges(vec![
+            Edge{from: a, to: b, token: t1, capacity: U256::from(10)},
+            Edge{from: a, to: c, token: t2, capacity: U256::from(7)},
+            Edge{from: b, to: d, token: t2, capacity: U256::from(9)},
+            Edge{from: c, to: d, token: t1, capacity: U256::from(8)},
+        ]);
+        let mut flow = compute_flow(&a, &d, &edges);
+        flow.1.sort();
+        assert_eq!(flow, (U256::from(16), vec![
+            Edge{from: a, to: b, token: t1, capacity: U256::from(9)},
+            Edge{from: a, to: c, token: t2, capacity: U256::from(7)},
+            Edge{from: b, to: d, token: t2, capacity: U256::from(9)},
+            Edge{from: c, to: d, token: t1, capacity: U256::from(7)},
+        ]));
+    }
+}
\ No newline at end of file
diff --git a/src/server.rs b/src/server.rs
index 71f6bbc..e05fcf8 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -69,7 +69,7 @@ fn handle_connection(
             );
             println!("Computed flow");
             // TODO error handling
-            socket.write_all(jsonrpc_result(request.id, json::JsonValue::from(flow)).as_bytes())?;
+            socket.write_all(jsonrpc_result(request.id, json::JsonValue::from(flow.0.to_string())).as_bytes())?;
         }
         "cancel" => {}
         "update_edges" => {}
diff --git a/src/types/address.rs b/src/types/address.rs
index 6d7a7d8..0ce85f6 100644
--- a/src/types/address.rs
+++ b/src/types/address.rs
@@ -1,8 +1,14 @@
-use std::fmt::{Display, Formatter};
+use std::fmt::{Display, Formatter, Debug};
 
-#[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq, PartialOrd)]
+#[derive(Clone, Copy, Default, Hash, Eq, PartialEq, Ord, PartialOrd)]
 pub struct Address([u8; 20]);
 
+impl Debug for Address {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "{}", self)
+    }
+}
+
 impl From<[u8; 20]> for Address {
     fn from(item: [u8; 20]) -> Self {
         Address(item)
diff --git a/src/types/edge.rs b/src/types/edge.rs
index e024dae..d8c6985 100644
--- a/src/types/edge.rs
+++ b/src/types/edge.rs
@@ -1,7 +1,7 @@
 use crate::types::Address;
 use crate::types::U256;
 
-#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, PartialOrd)]
+#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Ord, PartialOrd)]
 pub struct Edge {
     pub from: Address,
     pub to: Address,