From 9b3a0eecf86fd3e3c4906bd315ae2ba4b3955052 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Mon, 4 Mar 2024 18:11:36 +0000
Subject: [PATCH 01/13] Price improvement tests

---
 .../driver/src/tests/cases/protocol_fees.rs   | 129 ++++++++++++++++++
 crates/driver/src/tests/setup/mod.rs          |  30 ++++
 2 files changed, 159 insertions(+)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 42e526e32f..57d6187e54 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -9,6 +9,7 @@ use crate::{
             ab_solution,
             ExpectedOrderAmounts,
             FeePolicy,
+            PriceImprovementQuote,
             Test,
         },
     },
@@ -187,3 +188,131 @@ async fn volume_protocol_fee_sell_order() {
 
     protocol_fee_test_case(test_case).await;
 }
+
+#[tokio::test]
+#[ignore]
+async fn price_improvement_fee_buy_out_market_order() {
+    let quote_sell_amount = 21000000000000000000u128;
+    let quote_buy_amount = 18000000000000000000u128;
+    let fee_policy = FeePolicy::PriceImprovement {
+        factor: 0.5,
+        // high enough so we don't get capped by volume fee
+        max_volume_factor: 1.0,
+        quote: PriceImprovementQuote {
+            sell_amount: quote_sell_amount.into(),
+            buy_amount: quote_buy_amount.into(),
+            fee: 1000000000000000000u128.into(),
+        },
+    };
+    let executed_buy = 17143028023069342830u128;
+    let test_case = TestCase {
+        order_side: order::Side::Buy,
+        fee_policy,
+        order_sell_amount: 20000000000000000000u128.into(),
+        solver_fee: Some(10000000000000000000u128.into()),
+        quote_sell_amount: quote_sell_amount.into(),
+        quote_buy_amount: quote_buy_amount.into(),
+        executed: executed_buy.into(),
+        executed_sell_amount: 20476294902986820618u128.into(),
+        // executed buy amount should match order buy amount
+        executed_buy_amount: executed_buy.into(),
+    };
+
+    protocol_fee_test_case(test_case).await;
+}
+
+#[tokio::test]
+#[ignore]
+async fn price_improvement_fee_sell_out_market_order() {
+    let quote_sell_amount = 21000000000000000000u128;
+    let quote_buy_amount = 18000000000000000000u128;
+    let fee_policy = FeePolicy::PriceImprovement {
+        factor: 0.5,
+        // high enough so we don't get capped by volume fee
+        max_volume_factor: 1.0,
+        quote: PriceImprovementQuote {
+            sell_amount: quote_sell_amount.into(),
+            buy_amount: quote_buy_amount.into(),
+            fee: 1000000000000000000u128.into(),
+        },
+    };
+    let order_sell_amount = 20000000000000000000u128;
+    let test_case = TestCase {
+        order_side: order::Side::Sell,
+        fee_policy,
+        order_sell_amount: order_sell_amount.into(),
+        solver_fee: Some(10000000000000000000u128.into()),
+        quote_sell_amount: quote_sell_amount.into(),
+        quote_buy_amount: quote_buy_amount.into(),
+        executed: 10000000000000000000u128.into(),
+        // executed sell amount should match order sell amount
+        executed_sell_amount: order_sell_amount.into(),
+        executed_buy_amount: 16753332193352853234u128.into(),
+    };
+
+    protocol_fee_test_case(test_case).await;
+}
+
+#[tokio::test]
+#[ignore]
+async fn price_improvement_fee_buy_in_market_order() {
+    let quote_sell_amount = 17000000000000000000u128;
+    let quote_buy_amount = 10000000000000000000u128;
+    let fee_policy = FeePolicy::PriceImprovement {
+        factor: 0.5,
+        // high enough so we don't get capped by volume fee
+        max_volume_factor: 1.0,
+        quote: PriceImprovementQuote {
+            sell_amount: quote_sell_amount.into(),
+            buy_amount: quote_buy_amount.into(),
+            fee: 1000000000000000000u128.into(),
+        },
+    };
+    let executed_buy_amount = 11764354070151352996u128;
+    let test_case = TestCase {
+        order_side: order::Side::Buy,
+        fee_policy,
+        order_sell_amount: 20000000000000000000u128.into(),
+        solver_fee: Some(10000000000000000000u128.into()),
+        quote_sell_amount: quote_sell_amount.into(),
+        quote_buy_amount: quote_buy_amount.into(),
+        executed: executed_buy_amount.into(),
+        executed_sell_amount: 20587918663136217696u128.into(),
+        // executed buy amount should match order buy amount
+        executed_buy_amount: executed_buy_amount.into(),
+    };
+
+    protocol_fee_test_case(test_case).await;
+}
+
+#[tokio::test]
+#[ignore]
+async fn price_improvement_fee_sell_in_market_order() {
+    let quote_sell_amount = 9000000000000000000u128;
+    let quote_buy_amount = 25000000000000000000u128;
+    let fee_policy = FeePolicy::PriceImprovement {
+        factor: 0.5,
+        // high enough so we don't get capped by volume fee
+        max_volume_factor: 1.0,
+        quote: PriceImprovementQuote {
+            sell_amount: quote_sell_amount.into(),
+            buy_amount: quote_buy_amount.into(),
+            fee: 1000000000000000000u128.into(),
+        },
+    };
+    let order_sell_amount = 10000000000000000000u128;
+    let test_case = TestCase {
+        order_side: order::Side::Sell,
+        fee_policy,
+        order_sell_amount: order_sell_amount.into(),
+        solver_fee: Some(5000000000000000000u128.into()),
+        quote_sell_amount: quote_sell_amount.into(),
+        quote_buy_amount: quote_buy_amount.into(),
+        executed: 5000000000000000000u128.into(),
+        // executed sell amount should match order sell amount
+        executed_sell_amount: order_sell_amount.into(),
+        executed_buy_amount: 26388750430470970935u128.into(),
+    };
+
+    protocol_fee_test_case(test_case).await;
+}
diff --git a/crates/driver/src/tests/setup/mod.rs b/crates/driver/src/tests/setup/mod.rs
index b5bf85ef90..3e203406f5 100644
--- a/crates/driver/src/tests/setup/mod.rs
+++ b/crates/driver/src/tests/setup/mod.rs
@@ -72,6 +72,15 @@ pub enum Score {
     RiskAdjusted { success_probability: f64 },
 }
 
+#[serde_as]
+#[derive(Debug, Clone, PartialEq, serde::Serialize)]
+#[serde(rename_all = "camelCase", tag = "kind")]
+pub struct PriceImprovementQuote {
+    pub buy_amount: eth::U256,
+    pub sell_amount: eth::U256,
+    pub fee: eth::U256,
+}
+
 #[serde_as]
 #[derive(Debug, Clone, PartialEq, serde::Serialize)]
 #[serde(rename_all = "camelCase", tag = "kind")]
@@ -79,6 +88,12 @@ pub enum FeePolicy {
     #[serde(rename_all = "camelCase")]
     Surplus { factor: f64, max_volume_factor: f64 },
     #[serde(rename_all = "camelCase")]
+    PriceImprovement {
+        factor: f64,
+        max_volume_factor: f64,
+        quote: PriceImprovementQuote,
+    },
+    #[serde(rename_all = "camelCase")]
     Volume { factor: f64 },
 }
 
@@ -94,6 +109,21 @@ impl FeePolicy {
                     "maxVolumeFactor": max_volume_factor
                 }
             }),
+            FeePolicy::PriceImprovement {
+                factor,
+                max_volume_factor,
+                quote,
+            } => json!({
+                "priceImprovement": {
+                    "factor": factor,
+                    "maxVolumeFactor": max_volume_factor,
+                    "quote": {
+                        "sellAmount": quote.sell_amount,
+                        "buyAmount": quote.buy_amount,
+                        "fee": quote.fee,
+                    }
+                }
+            }),
             FeePolicy::Volume { factor } => json!({
                 "volume": {
                     "factor": factor

From dd233e23082f11509f1b1902004ac30b4b61e3bf Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Tue, 5 Mar 2024 11:01:39 +0000
Subject: [PATCH 02/13] Adjusted values

---
 .../driver/src/tests/cases/protocol_fees.rs   | 55 +++++++++----------
 1 file changed, 27 insertions(+), 28 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 57d6187e54..b6a91a1c38 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -192,8 +192,8 @@ async fn volume_protocol_fee_sell_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_out_market_order() {
-    let quote_sell_amount = 21000000000000000000u128;
-    let quote_buy_amount = 18000000000000000000u128;
+    let quote_sell_amount = 50000000000000000000u128;
+    let quote_buy_amount = 45000000000000000000u128;
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         // high enough so we don't get capped by volume fee
@@ -204,18 +204,18 @@ async fn price_improvement_fee_buy_out_market_order() {
             fee: 1000000000000000000u128.into(),
         },
     };
-    let executed_buy = 17143028023069342830u128;
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: 20000000000000000000u128.into(),
+        order_sell_amount: 50000000000000000000u128.into(),
         solver_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: quote_sell_amount.into(),
         quote_buy_amount: quote_buy_amount.into(),
-        executed: executed_buy.into(),
-        executed_sell_amount: 20476294902986820618u128.into(),
-        // executed buy amount should match order buy amount
-        executed_buy_amount: executed_buy.into(),
+        executed: quote_buy_amount.into(),
+        // sell amount + quote fee * factor
+        executed_sell_amount: 50500000000000000000u128.into(),
+        // executed buy amount should match quote buy amount
+        executed_buy_amount: quote_buy_amount.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -224,8 +224,8 @@ async fn price_improvement_fee_buy_out_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_out_market_order() {
-    let quote_sell_amount = 21000000000000000000u128;
-    let quote_buy_amount = 18000000000000000000u128;
+    let quote_sell_amount = 50000000000000000000u128;
+    let quote_buy_amount = 45000000000000000000u128;
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         // high enough so we don't get capped by volume fee
@@ -236,7 +236,7 @@ async fn price_improvement_fee_sell_out_market_order() {
             fee: 1000000000000000000u128.into(),
         },
     };
-    let order_sell_amount = 20000000000000000000u128;
+    let order_sell_amount = 50000000000000000000u128;
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
@@ -244,10 +244,10 @@ async fn price_improvement_fee_sell_out_market_order() {
         solver_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: quote_sell_amount.into(),
         quote_buy_amount: quote_buy_amount.into(),
-        executed: 10000000000000000000u128.into(),
+        executed: 40000000000000000000u128.into(),
         // executed sell amount should match order sell amount
         executed_sell_amount: order_sell_amount.into(),
-        executed_buy_amount: 16753332193352853234u128.into(),
+        executed_buy_amount: 44558823529411764706u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -256,8 +256,8 @@ async fn price_improvement_fee_sell_out_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_in_market_order() {
-    let quote_sell_amount = 17000000000000000000u128;
-    let quote_buy_amount = 10000000000000000000u128;
+    let quote_sell_amount = 50000000000000000000u128;
+    let quote_buy_amount = 45000000000000000000u128;
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         // high enough so we don't get capped by volume fee
@@ -265,21 +265,20 @@ async fn price_improvement_fee_buy_in_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: quote_sell_amount.into(),
             buy_amount: quote_buy_amount.into(),
-            fee: 1000000000000000000u128.into(),
+            fee: 20000000000000000000u128.into(),
         },
     };
-    let executed_buy_amount = 11764354070151352996u128;
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: 20000000000000000000u128.into(),
+        order_sell_amount: 50000000000000000000u128.into(),
         solver_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: quote_sell_amount.into(),
         quote_buy_amount: quote_buy_amount.into(),
-        executed: executed_buy_amount.into(),
-        executed_sell_amount: 20587918663136217696u128.into(),
+        executed: quote_buy_amount.into(),
+        executed_sell_amount: 60000000000000000000u128.into(),
         // executed buy amount should match order buy amount
-        executed_buy_amount: executed_buy_amount.into(),
+        executed_buy_amount: quote_buy_amount.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -288,8 +287,8 @@ async fn price_improvement_fee_buy_in_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_in_market_order() {
-    let quote_sell_amount = 9000000000000000000u128;
-    let quote_buy_amount = 25000000000000000000u128;
+    let quote_sell_amount = 50000000000000000000u128;
+    let quote_buy_amount = 45000000000000000000u128;
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         // high enough so we don't get capped by volume fee
@@ -297,21 +296,21 @@ async fn price_improvement_fee_sell_in_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: quote_sell_amount.into(),
             buy_amount: quote_buy_amount.into(),
-            fee: 1000000000000000000u128.into(),
+            fee: 20000000000000000000u128.into(),
         },
     };
-    let order_sell_amount = 10000000000000000000u128;
+    let order_sell_amount = 50000000000000000000u128;
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount: order_sell_amount.into(),
-        solver_fee: Some(5000000000000000000u128.into()),
+        solver_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: quote_sell_amount.into(),
         quote_buy_amount: quote_buy_amount.into(),
-        executed: 5000000000000000000u128.into(),
+        executed: 40000000000000000000u128.into(),
         // executed sell amount should match order sell amount
         executed_sell_amount: order_sell_amount.into(),
-        executed_buy_amount: 26388750430470970935u128.into(),
+        executed_buy_amount: 38571428571428571429u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;

From 3f7bf5e5a8bf16fc0218b4b1951a764f38de56ba Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Tue, 5 Mar 2024 17:28:19 +0000
Subject: [PATCH 03/13] Update crates/driver/src/tests/cases/protocol_fees.rs

Co-authored-by: Felix Leupold <felixleupold90@gmail.com>
---
 crates/driver/src/tests/cases/protocol_fees.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index b6a91a1c38..31734177d7 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -191,7 +191,7 @@ async fn volume_protocol_fee_sell_order() {
 
 #[tokio::test]
 #[ignore]
-async fn price_improvement_fee_buy_out_market_order() {
+async fn price_improvement_fee_buy_out_of_market_order() {
     let quote_sell_amount = 50000000000000000000u128;
     let quote_buy_amount = 45000000000000000000u128;
     let fee_policy = FeePolicy::PriceImprovement {

From 4b70eb70a3a5838508d3d82de076d99069c309b4 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Tue, 5 Mar 2024 17:28:24 +0000
Subject: [PATCH 04/13] Update crates/driver/src/tests/cases/protocol_fees.rs

Co-authored-by: Felix Leupold <felixleupold90@gmail.com>
---
 crates/driver/src/tests/cases/protocol_fees.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 31734177d7..9a7e6e7c90 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -223,7 +223,7 @@ async fn price_improvement_fee_buy_out_of_market_order() {
 
 #[tokio::test]
 #[ignore]
-async fn price_improvement_fee_sell_out_market_order() {
+async fn price_improvement_fee_sell_out_of_market_order() {
     let quote_sell_amount = 50000000000000000000u128;
     let quote_buy_amount = 45000000000000000000u128;
     let fee_policy = FeePolicy::PriceImprovement {

From 38135e86b9d40d46b22d99e6c4650d707fd4f479 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Tue, 5 Mar 2024 20:45:35 +0000
Subject: [PATCH 05/13] Adjusted values

---
 .../driver/src/tests/cases/protocol_fees.rs   | 119 +++++++++---------
 1 file changed, 60 insertions(+), 59 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 9a7e6e7c90..f440faf514 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -192,30 +192,30 @@ async fn volume_protocol_fee_sell_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_out_of_market_order() {
-    let quote_sell_amount = 50000000000000000000u128;
-    let quote_buy_amount = 45000000000000000000u128;
+    let quote_sell_amount = 50000000000000000000u128.into();
+    let quote_buy_amount = 35000000000000000000u128.into();
+    let fee: eth::U256 = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
-        // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount.into(),
-            buy_amount: quote_buy_amount.into(),
-            fee: 1000000000000000000u128.into(),
+            sell_amount: quote_sell_amount,
+            buy_amount: quote_buy_amount,
+            fee,
         },
     };
+    let order_sell_amount = 50000000000000000000u128.into();
+    let order_buy_amount = 40000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: quote_sell_amount.into(),
-        quote_buy_amount: quote_buy_amount.into(),
-        executed: quote_buy_amount.into(),
-        // sell amount + quote fee * factor
-        executed_sell_amount: 50500000000000000000u128.into(),
-        // executed buy amount should match quote buy amount
-        executed_buy_amount: quote_buy_amount.into(),
+        order_sell_amount,
+        solver_fee: Some(fee),
+        quote_sell_amount: order_sell_amount,
+        quote_buy_amount: order_buy_amount,
+        executed: order_buy_amount,
+        executed_sell_amount: 54142857142857142857u128.into(),
+        executed_buy_amount: order_buy_amount,
     };
 
     protocol_fee_test_case(test_case).await;
@@ -224,30 +224,30 @@ async fn price_improvement_fee_buy_out_of_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_out_of_market_order() {
-    let quote_sell_amount = 50000000000000000000u128;
-    let quote_buy_amount = 45000000000000000000u128;
+    let quote_sell_amount = 50000000000000000000u128.into();
+    let quote_buy_amount = 35000000000000000000u128.into();
+    let fee: eth::U256 = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
-        // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount.into(),
-            buy_amount: quote_buy_amount.into(),
-            fee: 1000000000000000000u128.into(),
+            sell_amount: quote_sell_amount,
+            buy_amount: quote_buy_amount,
+            fee,
         },
     };
-    let order_sell_amount = 50000000000000000000u128;
+    let order_sell_amount = 50000000000000000000u128.into();
+    let order_buy_amount = 40000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
-        order_sell_amount: order_sell_amount.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: quote_sell_amount.into(),
-        quote_buy_amount: quote_buy_amount.into(),
-        executed: 40000000000000000000u128.into(),
-        // executed sell amount should match order sell amount
-        executed_sell_amount: order_sell_amount.into(),
-        executed_buy_amount: 44558823529411764706u128.into(),
+        order_sell_amount,
+        solver_fee: Some(fee),
+        quote_sell_amount: order_sell_amount,
+        quote_buy_amount: order_buy_amount,
+        executed: order_sell_amount - fee,
+        executed_sell_amount: order_sell_amount,
+        executed_buy_amount: 37156862745098039215u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -256,29 +256,30 @@ async fn price_improvement_fee_sell_out_of_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_in_market_order() {
-    let quote_sell_amount = 50000000000000000000u128;
-    let quote_buy_amount = 45000000000000000000u128;
+    let quote_sell_amount: eth::U256 = 50000000000000000000u128.into();
+    let quote_buy_amount: eth::U256 = 40000000000000000000u128.into();
+    let fee = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
-        // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount.into(),
-            buy_amount: quote_buy_amount.into(),
-            fee: 20000000000000000000u128.into(),
+            sell_amount: quote_sell_amount,
+            buy_amount: quote_buy_amount,
+            fee,
         },
     };
+    let order_sell_amount = 50000000000000000000u128.into();
+    let order_buy_amount = 35000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: quote_sell_amount.into(),
-        quote_buy_amount: quote_buy_amount.into(),
-        executed: quote_buy_amount.into(),
-        executed_sell_amount: 60000000000000000000u128.into(),
-        // executed buy amount should match order buy amount
-        executed_buy_amount: quote_buy_amount.into(),
+        order_sell_amount: quote_sell_amount,
+        solver_fee: Some(fee),
+        quote_sell_amount: order_sell_amount,
+        quote_buy_amount: order_buy_amount,
+        executed: order_buy_amount,
+        executed_sell_amount: order_sell_amount,
+        executed_buy_amount: order_buy_amount,
     };
 
     protocol_fee_test_case(test_case).await;
@@ -287,30 +288,30 @@ async fn price_improvement_fee_buy_in_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_in_market_order() {
-    let quote_sell_amount = 50000000000000000000u128;
-    let quote_buy_amount = 45000000000000000000u128;
+    let quote_sell_amount: eth::U256 = 50000000000000000000u128.into();
+    let quote_buy_amount: eth::U256 = 40000000000000000000u128.into();
+    let fee = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
-        // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount.into(),
-            buy_amount: quote_buy_amount.into(),
-            fee: 20000000000000000000u128.into(),
+            sell_amount: quote_sell_amount,
+            buy_amount: quote_buy_amount,
+            fee,
         },
     };
-    let order_sell_amount = 50000000000000000000u128;
+    let order_sell_amount: eth::U256 = 50000000000000000000u128.into();
+    let order_buy_amount: eth::U256 = 35000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
-        order_sell_amount: order_sell_amount.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: quote_sell_amount.into(),
-        quote_buy_amount: quote_buy_amount.into(),
-        executed: 40000000000000000000u128.into(),
-        // executed sell amount should match order sell amount
-        executed_sell_amount: order_sell_amount.into(),
-        executed_buy_amount: 38571428571428571429u128.into(),
+        order_sell_amount,
+        solver_fee: Some(fee),
+        quote_sell_amount: order_sell_amount,
+        quote_buy_amount: order_buy_amount,
+        executed: order_sell_amount - fee,
+        executed_sell_amount: order_sell_amount,
+        executed_buy_amount: order_buy_amount,
     };
 
     protocol_fee_test_case(test_case).await;

From 287c1f923e0d8388fc4432ce0d189a8e1cabb1aa Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Tue, 5 Mar 2024 20:58:55 +0000
Subject: [PATCH 06/13] Explicit fee amounts

---
 .../driver/src/tests/cases/protocol_fees.rs   | 68 ++++++++-----------
 1 file changed, 29 insertions(+), 39 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index f440faf514..c66ede2a78 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -19,7 +19,7 @@ struct TestCase {
     order_side: order::Side,
     fee_policy: FeePolicy,
     order_sell_amount: eth::U256,
-    solver_fee: Option<eth::U256>,
+    network_fee: Option<eth::U256>,
     quote_sell_amount: eth::U256,
     quote_buy_amount: eth::U256,
     executed: eth::U256,
@@ -44,7 +44,7 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         .kind(order::Kind::Limit)
         .sell_amount(test_case.order_sell_amount)
         .side(test_case.order_side)
-        .solver_fee(test_case.solver_fee)
+        .solver_fee(test_case.network_fee)
         .fee_policy(test_case.fee_policy)
         .executed(test_case.executed)
         .expected_amounts(expected_amounts);
@@ -71,7 +71,7 @@ async fn surplus_protocol_fee_buy_order_not_capped() {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: 50000000000000000000u128.into(),
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
@@ -94,7 +94,7 @@ async fn surplus_protocol_fee_sell_order_not_capped() {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: 50000000000000000000u128.into(),
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
@@ -117,7 +117,7 @@ async fn surplus_protocol_fee_buy_order_capped() {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: 50000000000000000000u128.into(),
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
@@ -140,7 +140,7 @@ async fn surplus_protocol_fee_sell_order_capped() {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: 50000000000000000000u128.into(),
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
@@ -159,7 +159,7 @@ async fn volume_protocol_fee_buy_order() {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: 50000000000000000000u128.into(),
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
@@ -178,7 +178,7 @@ async fn volume_protocol_fee_sell_order() {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount: 50000000000000000000u128.into(),
-        solver_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: 50000000000000000000u128.into(),
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
@@ -192,16 +192,13 @@ async fn volume_protocol_fee_sell_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_out_of_market_order() {
-    let quote_sell_amount = 50000000000000000000u128.into();
-    let quote_buy_amount = 35000000000000000000u128.into();
-    let fee: eth::U256 = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount,
-            buy_amount: quote_buy_amount,
-            fee,
+            sell_amount: 50000000000000000000u128.into(),
+            buy_amount: 35000000000000000000u128.into(),
+            fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount = 50000000000000000000u128.into();
@@ -210,7 +207,7 @@ async fn price_improvement_fee_buy_out_of_market_order() {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
-        solver_fee: Some(fee),
+        network_fee: Some(2000000000000000000u128.into()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
@@ -224,28 +221,26 @@ async fn price_improvement_fee_buy_out_of_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_out_of_market_order() {
-    let quote_sell_amount = 50000000000000000000u128.into();
-    let quote_buy_amount = 35000000000000000000u128.into();
-    let fee: eth::U256 = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount,
-            buy_amount: quote_buy_amount,
-            fee,
+            sell_amount: 50000000000000000000u128.into(),
+            buy_amount: 35000000000000000000u128.into(),
+            fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount = 50000000000000000000u128.into();
     let order_buy_amount = 40000000000000000000u128.into();
+    let network_fee = 2000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount,
-        solver_fee: Some(fee),
+        network_fee: Some(network_fee),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
-        executed: order_sell_amount - fee,
+        executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
         executed_buy_amount: 37156862745098039215u128.into(),
     };
@@ -256,16 +251,13 @@ async fn price_improvement_fee_sell_out_of_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_in_market_order() {
-    let quote_sell_amount: eth::U256 = 50000000000000000000u128.into();
-    let quote_buy_amount: eth::U256 = 40000000000000000000u128.into();
-    let fee = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount,
-            buy_amount: quote_buy_amount,
-            fee,
+            sell_amount: 50000000000000000000u128.into(),
+            buy_amount: 40000000000000000000u128.into(),
+            fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount = 50000000000000000000u128.into();
@@ -273,8 +265,8 @@ async fn price_improvement_fee_buy_in_market_order() {
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: quote_sell_amount,
-        solver_fee: Some(fee),
+        order_sell_amount,
+        network_fee: Some(2000000000000000000u128.into()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
@@ -288,28 +280,26 @@ async fn price_improvement_fee_buy_in_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_in_market_order() {
-    let quote_sell_amount: eth::U256 = 50000000000000000000u128.into();
-    let quote_buy_amount: eth::U256 = 40000000000000000000u128.into();
-    let fee = 1000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: quote_sell_amount,
-            buy_amount: quote_buy_amount,
-            fee,
+            sell_amount: 50000000000000000000u128.into(),
+            buy_amount: 40000000000000000000u128.into(),
+            fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount: eth::U256 = 50000000000000000000u128.into();
     let order_buy_amount: eth::U256 = 35000000000000000000u128.into();
+    let network_fee = 20000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount,
-        solver_fee: Some(fee),
+        network_fee: Some(network_fee),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
-        executed: order_sell_amount - fee,
+        executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
         executed_buy_amount: order_buy_amount,
     };

From 08c7027a9d5568fa33cb1ede32ca40bd17acc95e Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Tue, 5 Mar 2024 21:03:25 +0000
Subject: [PATCH 07/13] Network fee

---
 crates/driver/src/tests/cases/protocol_fees.rs | 8 ++++----
 crates/driver/src/tests/setup/mod.rs           | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index c66ede2a78..0be7607576 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -198,7 +198,7 @@ async fn price_improvement_fee_buy_out_of_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
             buy_amount: 35000000000000000000u128.into(),
-            fee: 1000000000000000000u128.into(),
+            network_fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount = 50000000000000000000u128.into();
@@ -227,7 +227,7 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
             buy_amount: 35000000000000000000u128.into(),
-            fee: 1000000000000000000u128.into(),
+            network_fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount = 50000000000000000000u128.into();
@@ -257,7 +257,7 @@ async fn price_improvement_fee_buy_in_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
             buy_amount: 40000000000000000000u128.into(),
-            fee: 1000000000000000000u128.into(),
+            network_fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount = 50000000000000000000u128.into();
@@ -286,7 +286,7 @@ async fn price_improvement_fee_sell_in_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
             buy_amount: 40000000000000000000u128.into(),
-            fee: 1000000000000000000u128.into(),
+            network_fee: 1000000000000000000u128.into(),
         },
     };
     let order_sell_amount: eth::U256 = 50000000000000000000u128.into();
diff --git a/crates/driver/src/tests/setup/mod.rs b/crates/driver/src/tests/setup/mod.rs
index 3e203406f5..80b44ccf67 100644
--- a/crates/driver/src/tests/setup/mod.rs
+++ b/crates/driver/src/tests/setup/mod.rs
@@ -78,7 +78,7 @@ pub enum Score {
 pub struct PriceImprovementQuote {
     pub buy_amount: eth::U256,
     pub sell_amount: eth::U256,
-    pub fee: eth::U256,
+    pub network_fee: eth::U256,
 }
 
 #[serde_as]
@@ -120,7 +120,7 @@ impl FeePolicy {
                     "quote": {
                         "sellAmount": quote.sell_amount,
                         "buyAmount": quote.buy_amount,
-                        "fee": quote.fee,
+                        "fee": quote.network_fee,
                     }
                 }
             }),

From 758ebde1ad7970861373a54d16bd4e2c54f79163 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Wed, 6 Mar 2024 12:50:16 +0000
Subject: [PATCH 08/13] Adjusted fees

---
 .../driver/src/tests/cases/protocol_fees.rs   | 44 ++++++++++++-------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 0be7607576..4399253f7a 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -198,20 +198,23 @@ async fn price_improvement_fee_buy_out_of_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
             buy_amount: 35000000000000000000u128.into(),
-            network_fee: 1000000000000000000u128.into(),
+            network_fee: 20000000000000000000u128.into(),
         },
     };
-    let order_sell_amount = 50000000000000000000u128.into();
-    let order_buy_amount = 40000000000000000000u128.into();
+    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
+    // (45 + 10) * 35 < (50 + 20) * 35
+    // 1925 < 2450
+    let order_sell_amount = 45000000000000000000u128.into();
+    let order_buy_amount = 35000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
-        network_fee: Some(2000000000000000000u128.into()),
+        network_fee: Some(10000000000000000000u128.into()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        executed_sell_amount: 54142857142857142857u128.into(),
+        executed_sell_amount: 57500000000000000000u128.into(),
         executed_buy_amount: order_buy_amount,
     };
 
@@ -227,12 +230,15 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
             buy_amount: 35000000000000000000u128.into(),
-            network_fee: 1000000000000000000u128.into(),
+            network_fee: 20000000000000000000u128.into(),
         },
     };
+    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
+    // (50 + 10) * 35 < (50 + 20) * 40
+    // 2100 < 2800
     let order_sell_amount = 50000000000000000000u128.into();
     let order_buy_amount = 40000000000000000000u128.into();
-    let network_fee = 2000000000000000000u128.into();
+    let network_fee = 10000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
@@ -242,7 +248,7 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
-        executed_buy_amount: 37156862745098039215u128.into(),
+        executed_buy_amount: 32500000000000000000u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -255,22 +261,25 @@ async fn price_improvement_fee_buy_in_market_order() {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: 50000000000000000000u128.into(),
-            buy_amount: 40000000000000000000u128.into(),
-            network_fee: 1000000000000000000u128.into(),
+            sell_amount: 45000000000000000000u128.into(),
+            buy_amount: 35000000000000000000u128.into(),
+            network_fee: 10000000000000000000u128.into(),
         },
     };
+    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
+    // (50 + 15) * 35 < (45 + 10) * 35
+    // 2275 < 1925
     let order_sell_amount = 50000000000000000000u128.into();
     let order_buy_amount = 35000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
-        network_fee: Some(2000000000000000000u128.into()),
+        network_fee: Some(15000000000000000000u128.into()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        executed_sell_amount: order_sell_amount,
+        executed_sell_amount: 52500000000000000000u128.into(),
         executed_buy_amount: order_buy_amount,
     };
 
@@ -286,12 +295,15 @@ async fn price_improvement_fee_sell_in_market_order() {
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
             buy_amount: 40000000000000000000u128.into(),
-            network_fee: 1000000000000000000u128.into(),
+            network_fee: 10000000000000000000u128.into(),
         },
     };
+    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
+    // (50 + 15) * 40 < (50 + 10) * 35
+    // 2600 < 2100
     let order_sell_amount: eth::U256 = 50000000000000000000u128.into();
     let order_buy_amount: eth::U256 = 35000000000000000000u128.into();
-    let network_fee = 20000000000000000000u128.into();
+    let network_fee = 15000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
@@ -301,7 +313,7 @@ async fn price_improvement_fee_sell_in_market_order() {
         quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
-        executed_buy_amount: order_buy_amount,
+        executed_buy_amount: 34166666666666666667u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;

From 3a80fe613421c79d9736907175a35af87b0616e3 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Wed, 6 Mar 2024 20:14:07 +0000
Subject: [PATCH 09/13] Adjusted values

---
 .../driver/src/tests/cases/protocol_fees.rs   | 115 +++++++++++-------
 crates/driver/src/tests/setup/mod.rs          |   7 ++
 2 files changed, 76 insertions(+), 46 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 4399253f7a..812650b410 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -15,6 +15,8 @@ use crate::{
     },
 };
 
+const DEFAULT_SURPLUS_FACTOR: u64 = 2;
+
 struct TestCase {
     order_side: order::Side,
     fee_policy: FeePolicy,
@@ -41,6 +43,7 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         buy: test_case.executed_buy_amount,
     };
     let order = ab_order()
+        .surplus(DEFAULT_SURPLUS_FACTOR.into())
         .kind(order::Kind::Limit)
         .sell_amount(test_case.order_sell_amount)
         .side(test_case.order_side)
@@ -67,16 +70,18 @@ async fn surplus_protocol_fee_buy_order_not_capped() {
         // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
     };
+    let order_buy_amount = 40000000000000000000u128.into();
+    let order_sell_amount = 50000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
+        order_sell_amount,
         network_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: 50000000000000000000u128.into(),
-        quote_buy_amount: 40000000000000000000u128.into(),
-        executed: 40000000000000000000u128.into(),
-        executed_sell_amount: 100000000000000000000u128.into(),
-        executed_buy_amount: 40000000000000000000u128.into(),
+        quote_sell_amount: order_sell_amount,
+        quote_buy_amount: order_buy_amount,
+        executed: order_buy_amount,
+        executed_sell_amount: 75000000000000000000u128.into(),
+        executed_buy_amount: order_buy_amount,
     };
 
     protocol_fee_test_case(test_case).await;
@@ -90,16 +95,18 @@ async fn surplus_protocol_fee_sell_order_not_capped() {
         // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
     };
+    let order_sell_amount = 50000000000000000000u128.into();
+    let order_buy_amount = 40000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
+        order_sell_amount,
         network_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: 50000000000000000000u128.into(),
-        quote_buy_amount: 40000000000000000000u128.into(),
-        executed: 40000000000000000000u128.into(),
-        executed_sell_amount: 50000000000000000000u128.into(),
-        executed_buy_amount: 20000000002000000000u128.into(),
+        quote_sell_amount: order_sell_amount,
+        quote_buy_amount: order_buy_amount,
+        executed: order_buy_amount,
+        executed_sell_amount: order_sell_amount,
+        executed_buy_amount: 30000000000000000000u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -189,23 +196,31 @@ async fn volume_protocol_fee_sell_order() {
     protocol_fee_test_case(test_case).await;
 }
 
+// Price Improvement policy fee tests.
+// Out of market order could be defined as:
+//   (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
+// In the following tests Limit orders are used only, where order fee is 0. The
+// amount values are adjusted to respect the definition.
+
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_out_of_market_order() {
+    let order_sell_amount = 50000000000000000000u128.into();
+    let order_buy_amount = 40000000000000000000u128.into();
+    let quote_network_fee = 20000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: 50000000000000000000u128.into(),
-            buy_amount: 35000000000000000000u128.into(),
-            network_fee: 20000000000000000000u128.into(),
+            sell_amount: order_sell_amount,
+            // Since order surplus factor is 2.0, order's buy amount becomes x0.5 from the original
+            // value. Quote's buy amount is selected to be equal to the adjusted order's
+            // buy amount.
+            buy_amount: order_buy_amount / DEFAULT_SURPLUS_FACTOR,
+            // Quote's fee is high enough to make the order's conditions better.
+            network_fee: quote_network_fee,
         },
     };
-    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
-    // (45 + 10) * 35 < (50 + 20) * 35
-    // 1925 < 2450
-    let order_sell_amount = 45000000000000000000u128.into();
-    let order_buy_amount = 35000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
@@ -214,7 +229,9 @@ async fn price_improvement_fee_buy_out_of_market_order() {
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        executed_sell_amount: 57500000000000000000u128.into(),
+        // Executed values should be the same as for the surplus policy fee.
+        // ???
+        executed_sell_amount: 75000000000000000000u128.into(),
         executed_buy_amount: order_buy_amount,
     };
 
@@ -224,20 +241,20 @@ async fn price_improvement_fee_buy_out_of_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_out_of_market_order() {
+    let order_sell_amount = 50000000000000000000u128.into();
+    let order_buy_amount = 40000000000000000000u128.into();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: 50000000000000000000u128.into(),
-            buy_amount: 35000000000000000000u128.into(),
+            // Since order surplus factor is 2.0, order's sell amount becomes x2 from the original
+            // value. Quote's sell amount is selected to be equal to the adjusted order's
+            // buy amount.
+            sell_amount: order_sell_amount * DEFAULT_SURPLUS_FACTOR,
+            buy_amount: order_buy_amount,
             network_fee: 20000000000000000000u128.into(),
         },
     };
-    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
-    // (50 + 10) * 35 < (50 + 20) * 40
-    // 2100 < 2800
-    let order_sell_amount = 50000000000000000000u128.into();
-    let order_buy_amount = 40000000000000000000u128.into();
     let network_fee = 10000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
@@ -248,7 +265,9 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
-        executed_buy_amount: 32500000000000000000u128.into(),
+        // executed values should be the same as for the surplus policy fee
+        // ???
+        executed_buy_amount: 30000000000000000000u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -261,25 +280,28 @@ async fn price_improvement_fee_buy_in_market_order() {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: 45000000000000000000u128.into(),
-            buy_amount: 35000000000000000000u128.into(),
-            network_fee: 10000000000000000000u128.into(),
+            sell_amount: 50000000000000000000u128.into(),
+            // Since order surplus factor is 2.0, order's buy amount will be x0.5.
+            // To make it inside market price, the quote buy amount should be higher than the
+            // adjusted order's value.
+            buy_amount: 40000000000000000000u128.into(),
+            network_fee: 20000000000000000000u128.into(),
         },
     };
-    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
-    // (50 + 15) * 35 < (45 + 10) * 35
-    // 2275 < 1925
     let order_sell_amount = 50000000000000000000u128.into();
-    let order_buy_amount = 35000000000000000000u128.into();
+    let order_buy_amount = 40000000000000000000u128.into();
+    let network_fee = 10000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
-        network_fee: Some(15000000000000000000u128.into()),
+        network_fee: Some(network_fee),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        executed_sell_amount: 52500000000000000000u128.into(),
+        // quote.sell_amount + factor * quote.network_fee
+        // ???
+        executed_sell_amount: 60000000000000000000u128.into(),
         executed_buy_amount: order_buy_amount,
     };
 
@@ -293,17 +315,17 @@ async fn price_improvement_fee_sell_in_market_order() {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: 50000000000000000000u128.into(),
+            // Since order surplus factor is 2.0, order's sell amount will be x2.
+            // To make it inside market price, the quote sell amount should be lower than the
+            // adjusted order's value.
+            sell_amount: 500000000000000000000u128.into(),
             buy_amount: 40000000000000000000u128.into(),
-            network_fee: 10000000000000000000u128.into(),
+            network_fee: 20000000000000000000u128.into(),
         },
     };
-    // (order.sell + order.fee) * quote.buy < (quote.sell + quote.fee) * order.buy
-    // (50 + 15) * 40 < (50 + 10) * 35
-    // 2600 < 2100
     let order_sell_amount: eth::U256 = 50000000000000000000u128.into();
-    let order_buy_amount: eth::U256 = 35000000000000000000u128.into();
-    let network_fee = 15000000000000000000u128.into();
+    let order_buy_amount: eth::U256 = 40000000000000000000u128.into();
+    let network_fee = 10000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
@@ -313,7 +335,8 @@ async fn price_improvement_fee_sell_in_market_order() {
         quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
-        executed_buy_amount: 34166666666666666667u128.into(),
+        // ???
+        executed_buy_amount: 30000000000000000000u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;
diff --git a/crates/driver/src/tests/setup/mod.rs b/crates/driver/src/tests/setup/mod.rs
index 80b44ccf67..6ad678fe1b 100644
--- a/crates/driver/src/tests/setup/mod.rs
+++ b/crates/driver/src/tests/setup/mod.rs
@@ -236,6 +236,13 @@ impl Order {
         }
     }
 
+    pub fn surplus(self, surplus_factor: eth::U256) -> Self {
+        Self {
+            surplus_factor,
+            ..self
+        }
+    }
+
     /// Mark this order as internalizable.
     pub fn internalize(self) -> Self {
         Self {

From 75d99c135f6411f2469f8dd146af89914d6a5c32 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Fri, 8 Mar 2024 16:49:45 +0000
Subject: [PATCH 10/13] Order surplus factor=1

---
 .../driver/src/tests/cases/protocol_fees.rs   | 54 +++++++++----------
 crates/driver/src/tests/setup/blockchain.rs   | 26 +++++++--
 2 files changed, 48 insertions(+), 32 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 812650b410..69a622c977 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -27,6 +27,7 @@ struct TestCase {
     executed: eth::U256,
     executed_sell_amount: eth::U256,
     executed_buy_amount: eth::U256,
+    surplus_factor: eth::U256,
 }
 
 async fn protocol_fee_test_case(test_case: TestCase) {
@@ -43,7 +44,7 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         buy: test_case.executed_buy_amount,
     };
     let order = ab_order()
-        .surplus(DEFAULT_SURPLUS_FACTOR.into())
+        .surplus(test_case.surplus_factor)
         .kind(order::Kind::Limit)
         .sell_amount(test_case.order_sell_amount)
         .side(test_case.order_side)
@@ -82,6 +83,7 @@ async fn surplus_protocol_fee_buy_order_not_capped() {
         executed: order_buy_amount,
         executed_sell_amount: 75000000000000000000u128.into(),
         executed_buy_amount: order_buy_amount,
+        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -107,6 +109,7 @@ async fn surplus_protocol_fee_sell_order_not_capped() {
         executed: order_buy_amount,
         executed_sell_amount: order_sell_amount,
         executed_buy_amount: 30000000000000000000u128.into(),
+        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -130,6 +133,7 @@ async fn surplus_protocol_fee_buy_order_capped() {
         executed: 40000000000000000000u128.into(),
         executed_sell_amount: 55000000000000000000u128.into(),
         executed_buy_amount: 40000000000000000000u128.into(),
+        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -153,6 +157,7 @@ async fn surplus_protocol_fee_sell_order_capped() {
         executed: 40000000000000000000u128.into(),
         executed_sell_amount: 50000000000000000000u128.into(),
         executed_buy_amount: 35000000000000000000u128.into(),
+        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -172,6 +177,7 @@ async fn volume_protocol_fee_buy_order() {
         executed: 40000000000000000000u128.into(),
         executed_sell_amount: 75000000000000000000u128.into(),
         executed_buy_amount: 40000000000000000000u128.into(),
+        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -191,6 +197,7 @@ async fn volume_protocol_fee_sell_order() {
         executed: 40000000000000000000u128.into(),
         executed_sell_amount: 50000000000000000000u128.into(),
         executed_buy_amount: 15000000000000000000u128.into(),
+        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -213,11 +220,7 @@ async fn price_improvement_fee_buy_out_of_market_order() {
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
             sell_amount: order_sell_amount,
-            // Since order surplus factor is 2.0, order's buy amount becomes x0.5 from the original
-            // value. Quote's buy amount is selected to be equal to the adjusted order's
-            // buy amount.
             buy_amount: order_buy_amount / DEFAULT_SURPLUS_FACTOR,
-            // Quote's fee is high enough to make the order's conditions better.
             network_fee: quote_network_fee,
         },
     };
@@ -225,14 +228,14 @@ async fn price_improvement_fee_buy_out_of_market_order() {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
-        network_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(1000000000000000000u128.into()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        // Executed values should be the same as for the surplus policy fee.
-        // ???
-        executed_sell_amount: 75000000000000000000u128.into(),
+        // order sell amount + quote network fee * factor
+        executed_sell_amount: 60000000000000000000u128.into(),
         executed_buy_amount: order_buy_amount,
+        surplus_factor: 1.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -247,9 +250,6 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            // Since order surplus factor is 2.0, order's sell amount becomes x2 from the original
-            // value. Quote's sell amount is selected to be equal to the adjusted order's
-            // buy amount.
             sell_amount: order_sell_amount * DEFAULT_SURPLUS_FACTOR,
             buy_amount: order_buy_amount,
             network_fee: 20000000000000000000u128.into(),
@@ -265,9 +265,9 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
-        // executed values should be the same as for the surplus policy fee
-        // ???
-        executed_buy_amount: 30000000000000000000u128.into(),
+        // order buy amount - quote network fee * factor (in sell token)
+        executed_buy_amount: 32000000000000000000u128.into(),
+        surplus_factor: 1.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -281,14 +281,11 @@ async fn price_improvement_fee_buy_in_market_order() {
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
             sell_amount: 50000000000000000000u128.into(),
-            // Since order surplus factor is 2.0, order's buy amount will be x0.5.
-            // To make it inside market price, the quote buy amount should be higher than the
-            // adjusted order's value.
             buy_amount: 40000000000000000000u128.into(),
             network_fee: 20000000000000000000u128.into(),
         },
     };
-    let order_sell_amount = 50000000000000000000u128.into();
+    let order_sell_amount = 100000000000000000000u128.into();
     let order_buy_amount = 40000000000000000000u128.into();
     let network_fee = 10000000000000000000u128.into();
     let test_case = TestCase {
@@ -299,10 +296,11 @@ async fn price_improvement_fee_buy_in_market_order() {
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        // quote.sell_amount + factor * quote.network_fee
-        // ???
-        executed_sell_amount: 60000000000000000000u128.into(),
+        // todo: no price improvement since quote provides better conditions, but the fee isn't
+        // considered
+        executed_sell_amount: order_sell_amount,
         executed_buy_amount: order_buy_amount,
+        surplus_factor: 1.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -315,16 +313,13 @@ async fn price_improvement_fee_sell_in_market_order() {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            // Since order surplus factor is 2.0, order's sell amount will be x2.
-            // To make it inside market price, the quote sell amount should be lower than the
-            // adjusted order's value.
             sell_amount: 500000000000000000000u128.into(),
             buy_amount: 40000000000000000000u128.into(),
             network_fee: 20000000000000000000u128.into(),
         },
     };
     let order_sell_amount: eth::U256 = 50000000000000000000u128.into();
-    let order_buy_amount: eth::U256 = 40000000000000000000u128.into();
+    let order_buy_amount: eth::U256 = 10000000000000000000u128.into();
     let network_fee = 10000000000000000000u128.into();
     let test_case = TestCase {
         order_side: order::Side::Sell,
@@ -334,9 +329,10 @@ async fn price_improvement_fee_sell_in_market_order() {
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
-        executed_sell_amount: order_sell_amount,
-        // ???
-        executed_buy_amount: 30000000000000000000u128.into(),
+        // todo: explain values
+        executed_sell_amount: 80000000000000000000u128.into(),
+        executed_buy_amount: 36000000000000000000u128.into(),
+        surplus_factor: 1.into(),
     };
 
     protocol_fee_test_case(test_case).await;
diff --git a/crates/driver/src/tests/setup/blockchain.rs b/crates/driver/src/tests/setup/blockchain.rs
index 864d622a85..820e3d17a5 100644
--- a/crates/driver/src/tests/setup/blockchain.rs
+++ b/crates/driver/src/tests/setup/blockchain.rs
@@ -6,7 +6,7 @@ use {
             eth::{self, ContractAddress},
         },
         infra::time,
-        tests::{self, boundary},
+        tests::{self, boundary, setup::FeePolicy},
     },
     ethcontract::{dyns::DynWeb3, transport::DynTransport, Web3},
     futures::Future,
@@ -97,14 +97,34 @@ impl QuotedOrder {
     pub fn buy_amount(&self) -> eth::U256 {
         match self.order.side {
             order::Side::Buy => self.buy,
-            order::Side::Sell => self.buy / self.order.surplus_factor,
+            order::Side::Sell => {
+                let fee = match &self.order.fee_policy {
+                    FeePolicy::PriceImprovement {
+                        factor: _,
+                        max_volume_factor: _,
+                        quote,
+                    } => quote.network_fee,
+                    _ => eth::U256::zero(),
+                };
+                (self.buy * (self.sell - fee) / self.sell) / self.order.surplus_factor
+            }
         }
     }
 
     /// The sell amount with the surplus factor.
     pub fn sell_amount(&self) -> eth::U256 {
         match self.order.side {
-            order::Side::Buy => self.sell * self.order.surplus_factor,
+            order::Side::Buy => {
+                let fee = match &self.order.fee_policy {
+                    FeePolicy::PriceImprovement {
+                        factor: _,
+                        max_volume_factor: _,
+                        quote,
+                    } => quote.network_fee,
+                    _ => eth::U256::zero(),
+                };
+                (self.sell + fee) * self.order.surplus_factor
+            }
             order::Side::Sell => self.sell,
         }
     }

From 1a473437e3538b7bc86b4981f064bb35cdf06ac9 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Fri, 8 Mar 2024 17:05:57 +0000
Subject: [PATCH 11/13] Fixed values

---
 .../driver/src/tests/cases/protocol_fees.rs   | 123 +++++++++---------
 1 file changed, 62 insertions(+), 61 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 71ca3b5824..f65a93cd73 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -72,17 +72,17 @@ async fn surplus_protocol_fee_buy_order_not_capped() {
         // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
     };
-    let order_buy_amount = 40000000000000000000u128.into();
-    let order_sell_amount = 50000000000000000000u128.into();
+    let order_buy_amount = 40.ether().into_wei();
+    let order_sell_amount = 50.ether().into_wei();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
-        network_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10.ether().into_wei()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        executed_sell_amount: 75000000000000000000u128.into(),
+        executed_sell_amount: 75.ether().into_wei(),
         executed_buy_amount: order_buy_amount,
         surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
@@ -98,18 +98,18 @@ async fn surplus_protocol_fee_sell_order_not_capped() {
         // high enough so we don't get capped by volume fee
         max_volume_factor: 1.0,
     };
-    let order_sell_amount = 50000000000000000000u128.into();
-    let order_buy_amount = 40000000000000000000u128.into();
+    let order_sell_amount = 50.ether().into_wei();
+    let order_buy_amount = 40.ether().into_wei();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount,
-        network_fee: Some(10000000000000000000u128.into()),
+        network_fee: Some(10.ether().into_wei()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
         executed_sell_amount: order_sell_amount,
-        executed_buy_amount: 30000000000000000000u128.into(),
+        executed_buy_amount: 30.ether().into_wei(),
         surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
@@ -127,13 +127,13 @@ async fn surplus_protocol_fee_buy_order_capped() {
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
-        network_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: 50000000000000000000u128.into(),
-        quote_buy_amount: 40000000000000000000u128.into(),
-        executed: 40000000000000000000u128.into(),
-        executed_sell_amount: 55000000000000000000u128.into(),
-        executed_buy_amount: 40000000000000000000u128.into(),
+        order_sell_amount: 50.ether().into_wei(),
+        network_fee: Some(10.ether().into_wei()),
+        quote_sell_amount: 50.ether().into_wei(),
+        quote_buy_amount: 40.ether().into_wei(),
+        executed: 40.ether().into_wei(),
+        executed_sell_amount: 55.ether().into_wei(),
+        executed_buy_amount: 40.ether().into_wei(),
         surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
@@ -151,13 +151,13 @@ async fn surplus_protocol_fee_sell_order_capped() {
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
-        network_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: 50000000000000000000u128.into(),
-        quote_buy_amount: 40000000000000000000u128.into(),
-        executed: 40000000000000000000u128.into(),
-        executed_sell_amount: 50000000000000000000u128.into(),
-        executed_buy_amount: 35000000000000000000u128.into(),
+        order_sell_amount: 50.ether().into_wei(),
+        network_fee: Some(10.ether().into_wei()),
+        quote_sell_amount: 50.ether().into_wei(),
+        quote_buy_amount: 40.ether().into_wei(),
+        executed: 40.ether().into_wei(),
+        executed_sell_amount: 50.ether().into_wei(),
+        executed_buy_amount: 36.ether().into_wei(),
         surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
@@ -171,13 +171,13 @@ async fn volume_protocol_fee_buy_order() {
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
-        network_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: 50000000000000000000u128.into(),
-        quote_buy_amount: 40000000000000000000u128.into(),
-        executed: 40000000000000000000u128.into(),
-        executed_sell_amount: 75000000000000000000u128.into(),
-        executed_buy_amount: 40000000000000000000u128.into(),
+        order_sell_amount: 50.ether().into_wei(),
+        network_fee: Some(10.ether().into_wei()),
+        quote_sell_amount: 50.ether().into_wei(),
+        quote_buy_amount: 40.ether().into_wei(),
+        executed: 40.ether().into_wei(),
+        executed_sell_amount: 75.ether().into_wei(),
+        executed_buy_amount: 40.ether().into_wei(),
         surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
@@ -191,13 +191,13 @@ async fn volume_protocol_fee_sell_order() {
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
-        order_sell_amount: 50000000000000000000u128.into(),
-        network_fee: Some(10000000000000000000u128.into()),
-        quote_sell_amount: 50000000000000000000u128.into(),
-        quote_buy_amount: 40000000000000000000u128.into(),
-        executed: 40000000000000000000u128.into(),
-        executed_sell_amount: 50000000000000000000u128.into(),
-        executed_buy_amount: 15000000000000000000u128.into(),
+        order_sell_amount: 50.ether().into_wei(),
+        network_fee: Some(10.ether().into_wei()),
+        quote_sell_amount: 50.ether().into_wei(),
+        quote_buy_amount: 40.ether().into_wei(),
+        executed: 40.ether().into_wei(),
+        executed_sell_amount: 50.ether().into_wei(),
+        executed_buy_amount: 20.ether().into_wei(),
         surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
@@ -213,9 +213,9 @@ async fn volume_protocol_fee_sell_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_out_of_market_order() {
-    let order_sell_amount = 50000000000000000000u128.into();
-    let order_buy_amount = 40000000000000000000u128.into();
-    let quote_network_fee = 20000000000000000000u128.into();
+    let order_sell_amount = 50.ether().into_wei();
+    let order_buy_amount = 40.ether().into_wei();
+    let quote_network_fee = 20.ether().into_wei();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
@@ -229,12 +229,12 @@ async fn price_improvement_fee_buy_out_of_market_order() {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
-        network_fee: Some(1000000000000000000u128.into()),
+        network_fee: Some(1.ether().into_wei()),
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
         // order sell amount + quote network fee * factor
-        executed_sell_amount: 60000000000000000000u128.into(),
+        executed_sell_amount: 60.ether().into_wei(),
         executed_buy_amount: order_buy_amount,
         surplus_factor: 1.into(),
     };
@@ -245,18 +245,18 @@ async fn price_improvement_fee_buy_out_of_market_order() {
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_out_of_market_order() {
-    let order_sell_amount = 50000000000000000000u128.into();
-    let order_buy_amount = 40000000000000000000u128.into();
+    let order_sell_amount = 50.ether().into_wei();
+    let order_buy_amount = 40.ether().into_wei();
     let fee_policy = FeePolicy::PriceImprovement {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
             sell_amount: order_sell_amount * DEFAULT_SURPLUS_FACTOR,
             buy_amount: order_buy_amount,
-            network_fee: 20000000000000000000u128.into(),
+            network_fee: 20.ether().into_wei(),
         },
     };
-    let network_fee = 10000000000000000000u128.into();
+    let network_fee = 10.ether().into_wei();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
@@ -267,7 +267,7 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
         // order buy amount - quote network fee * factor (in sell token)
-        executed_buy_amount: 32000000000000000000u128.into(),
+        executed_buy_amount: 32.ether().into_wei(),
         surplus_factor: 1.into(),
     };
 
@@ -281,14 +281,14 @@ async fn price_improvement_fee_buy_in_market_order() {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: 50000000000000000000u128.into(),
-            buy_amount: 40000000000000000000u128.into(),
-            network_fee: 20000000000000000000u128.into(),
+            sell_amount: 50.ether().into_wei(),
+            buy_amount: 40.ether().into_wei(),
+            network_fee: 20.ether().into_wei(),
         },
     };
-    let order_sell_amount = 100000000000000000000u128.into();
-    let order_buy_amount = 40000000000000000000u128.into();
-    let network_fee = 10000000000000000000u128.into();
+    let order_sell_amount = 100.ether().into_wei();
+    let order_buy_amount = 40.ether().into_wei();
+    let network_fee = 10.ether().into_wei();
     let test_case = TestCase {
         order_side: order::Side::Buy,
         fee_policy,
@@ -314,14 +314,14 @@ async fn price_improvement_fee_sell_in_market_order() {
         factor: 0.5,
         max_volume_factor: 1.0,
         quote: PriceImprovementQuote {
-            sell_amount: 500000000000000000000u128.into(),
-            buy_amount: 40000000000000000000u128.into(),
-            network_fee: 20000000000000000000u128.into(),
+            sell_amount: 50.ether().into_wei(),
+            buy_amount: 40.ether().into_wei(),
+            network_fee: 20.ether().into_wei(),
         },
     };
-    let order_sell_amount: eth::U256 = 50000000000000000000u128.into();
-    let order_buy_amount: eth::U256 = 10000000000000000000u128.into();
-    let network_fee = 10000000000000000000u128.into();
+    let order_sell_amount: eth::U256 = 50.ether().into_wei();
+    let order_buy_amount: eth::U256 = 10.ether().into_wei();
+    let network_fee = 10.ether().into_wei();
     let test_case = TestCase {
         order_side: order::Side::Sell,
         fee_policy,
@@ -330,9 +330,10 @@ async fn price_improvement_fee_sell_in_market_order() {
         quote_sell_amount: order_sell_amount,
         quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
-        // todo: explain values
-        executed_sell_amount: 80000000000000000000u128.into(),
-        executed_buy_amount: 36000000000000000000u128.into(),
+        // todo: no price improvement since quote provides better conditions, but the fee isn't
+        // considered
+        executed_sell_amount: 50.ether().into_wei(),
+        executed_buy_amount: 10.ether().into_wei(),
         surplus_factor: 1.into(),
     };
 

From 9daf971a5836e91a38b722fc20e6ce56d060546d Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Fri, 8 Mar 2024 18:17:30 +0000
Subject: [PATCH 12/13] Refactoring

---
 .../driver/src/tests/cases/protocol_fees.rs   | 199 ++++++++++--------
 1 file changed, 111 insertions(+), 88 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index f65a93cd73..a61fb13382 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -18,7 +18,7 @@ use crate::{
 
 const DEFAULT_SURPLUS_FACTOR: u64 = 2;
 
-struct TestCase {
+struct CommonTestCase {
     order_side: order::Side,
     fee_policy: FeePolicy,
     order_sell_amount: eth::U256,
@@ -28,10 +28,22 @@ struct TestCase {
     executed: eth::U256,
     executed_sell_amount: eth::U256,
     executed_buy_amount: eth::U256,
-    surplus_factor: eth::U256,
 }
 
-async fn protocol_fee_test_case(test_case: TestCase) {
+struct PriceImprovementTestCase {
+    order_side: order::Side,
+    policy_factor: f64,
+    policy_max_volume_factor: f64,
+    quote: PriceImprovementQuote,
+    order_sell_amount: eth::U256,
+    order_buy_amount: eth::U256,
+    network_fee: Option<eth::U256>,
+    executed: eth::U256,
+    executed_sell_amount: eth::U256,
+    executed_buy_amount: eth::U256,
+}
+
+async fn common_fee_test_case(test_case: CommonTestCase) {
     let test_name = format!(
         "Protocol Fee: {:?} {:?}",
         test_case.order_side, test_case.fee_policy
@@ -45,7 +57,7 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         buy: test_case.executed_buy_amount,
     };
     let order = ab_order()
-        .surplus(test_case.surplus_factor)
+        .surplus(DEFAULT_SURPLUS_FACTOR.into())
         .kind(order::Kind::Limit)
         .sell_amount(test_case.order_sell_amount)
         .side(test_case.order_side)
@@ -64,6 +76,41 @@ async fn protocol_fee_test_case(test_case: TestCase) {
     test.solve().await.ok().orders(&[order]);
 }
 
+async fn price_improvement_fee_test_case(test_case: PriceImprovementTestCase) {
+    let test_name = format!("Protocol Fee: {:?} PriceImprovement", test_case.order_side);
+    let liquidity_quote = ab_liquidity_quote()
+        .sell_amount(test_case.order_sell_amount)
+        .buy_amount(test_case.order_buy_amount);
+    let pool = ab_adjusted_pool(liquidity_quote);
+    let expected_amounts = ExpectedOrderAmounts {
+        sell: test_case.executed_sell_amount,
+        buy: test_case.executed_buy_amount,
+    };
+    let fee_policy = FeePolicy::PriceImprovement {
+        factor: test_case.policy_factor,
+        max_volume_factor: test_case.policy_max_volume_factor,
+        quote: test_case.quote,
+    };
+    let order = ab_order()
+        .no_surplus()
+        .kind(order::Kind::Limit)
+        .sell_amount(test_case.order_sell_amount)
+        .side(test_case.order_side)
+        .solver_fee(test_case.network_fee)
+        .fee_policy(fee_policy)
+        .executed(test_case.executed)
+        .expected_amounts(expected_amounts);
+    let test: Test = tests::setup()
+        .name(test_name)
+        .pool(pool)
+        .order(order.clone())
+        .solution(ab_solution())
+        .done()
+        .await;
+
+    test.solve().await.ok().orders(&[order]);
+}
+
 #[tokio::test]
 #[ignore]
 async fn surplus_protocol_fee_buy_order_not_capped() {
@@ -74,7 +121,7 @@ async fn surplus_protocol_fee_buy_order_not_capped() {
     };
     let order_buy_amount = 40.ether().into_wei();
     let order_sell_amount = 50.ether().into_wei();
-    let test_case = TestCase {
+    let test_case = CommonTestCase {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount,
@@ -84,10 +131,9 @@ async fn surplus_protocol_fee_buy_order_not_capped() {
         executed: order_buy_amount,
         executed_sell_amount: 75.ether().into_wei(),
         executed_buy_amount: order_buy_amount,
-        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    common_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
@@ -100,7 +146,7 @@ async fn surplus_protocol_fee_sell_order_not_capped() {
     };
     let order_sell_amount = 50.ether().into_wei();
     let order_buy_amount = 40.ether().into_wei();
-    let test_case = TestCase {
+    let test_case = CommonTestCase {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount,
@@ -110,10 +156,9 @@ async fn surplus_protocol_fee_sell_order_not_capped() {
         executed: order_buy_amount,
         executed_sell_amount: order_sell_amount,
         executed_buy_amount: 30.ether().into_wei(),
-        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    common_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
@@ -124,7 +169,7 @@ async fn surplus_protocol_fee_buy_order_capped() {
         // low enough so we get capped by volume fee
         max_volume_factor: 0.1,
     };
-    let test_case = TestCase {
+    let test_case = CommonTestCase {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount: 50.ether().into_wei(),
@@ -134,10 +179,9 @@ async fn surplus_protocol_fee_buy_order_capped() {
         executed: 40.ether().into_wei(),
         executed_sell_amount: 55.ether().into_wei(),
         executed_buy_amount: 40.ether().into_wei(),
-        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    common_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
@@ -148,7 +192,7 @@ async fn surplus_protocol_fee_sell_order_capped() {
         // low enough so we get capped by volume fee
         max_volume_factor: 0.1,
     };
-    let test_case = TestCase {
+    let test_case = CommonTestCase {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount: 50.ether().into_wei(),
@@ -158,17 +202,16 @@ async fn surplus_protocol_fee_sell_order_capped() {
         executed: 40.ether().into_wei(),
         executed_sell_amount: 50.ether().into_wei(),
         executed_buy_amount: 36.ether().into_wei(),
-        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    common_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
 #[ignore]
 async fn volume_protocol_fee_buy_order() {
     let fee_policy = FeePolicy::Volume { factor: 0.5 };
-    let test_case = TestCase {
+    let test_case = CommonTestCase {
         order_side: order::Side::Buy,
         fee_policy,
         order_sell_amount: 50.ether().into_wei(),
@@ -178,17 +221,16 @@ async fn volume_protocol_fee_buy_order() {
         executed: 40.ether().into_wei(),
         executed_sell_amount: 75.ether().into_wei(),
         executed_buy_amount: 40.ether().into_wei(),
-        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    common_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
 #[ignore]
 async fn volume_protocol_fee_sell_order() {
     let fee_policy = FeePolicy::Volume { factor: 0.5 };
-    let test_case = TestCase {
+    let test_case = CommonTestCase {
         order_side: order::Side::Sell,
         fee_policy,
         order_sell_amount: 50.ether().into_wei(),
@@ -198,10 +240,9 @@ async fn volume_protocol_fee_sell_order() {
         executed: 40.ether().into_wei(),
         executed_sell_amount: 50.ether().into_wei(),
         executed_buy_amount: 20.ether().into_wei(),
-        surplus_factor: DEFAULT_SURPLUS_FACTOR.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    common_fee_test_case(test_case).await;
 }
 
 // Price Improvement policy fee tests.
@@ -216,30 +257,26 @@ async fn price_improvement_fee_buy_out_of_market_order() {
     let order_sell_amount = 50.ether().into_wei();
     let order_buy_amount = 40.ether().into_wei();
     let quote_network_fee = 20.ether().into_wei();
-    let fee_policy = FeePolicy::PriceImprovement {
-        factor: 0.5,
-        max_volume_factor: 1.0,
-        quote: PriceImprovementQuote {
-            sell_amount: order_sell_amount,
-            buy_amount: order_buy_amount / DEFAULT_SURPLUS_FACTOR,
-            network_fee: quote_network_fee,
-        },
+    let quote = PriceImprovementQuote {
+        sell_amount: order_sell_amount,
+        buy_amount: order_buy_amount,
+        network_fee: quote_network_fee,
     };
-    let test_case = TestCase {
+    let test_case = PriceImprovementTestCase {
         order_side: order::Side::Buy,
-        fee_policy,
+        policy_factor: 0.5,
+        policy_max_volume_factor: 1.0,
+        quote,
         order_sell_amount,
+        order_buy_amount,
         network_fee: Some(1.ether().into_wei()),
-        quote_sell_amount: order_sell_amount,
-        quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
         // order sell amount + quote network fee * factor
         executed_sell_amount: 60.ether().into_wei(),
         executed_buy_amount: order_buy_amount,
-        surplus_factor: 1.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    price_improvement_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
@@ -247,95 +284,81 @@ async fn price_improvement_fee_buy_out_of_market_order() {
 async fn price_improvement_fee_sell_out_of_market_order() {
     let order_sell_amount = 50.ether().into_wei();
     let order_buy_amount = 40.ether().into_wei();
-    let fee_policy = FeePolicy::PriceImprovement {
-        factor: 0.5,
-        max_volume_factor: 1.0,
-        quote: PriceImprovementQuote {
-            sell_amount: order_sell_amount * DEFAULT_SURPLUS_FACTOR,
-            buy_amount: order_buy_amount,
-            network_fee: 20.ether().into_wei(),
-        },
+    let quote = PriceImprovementQuote {
+        sell_amount: order_sell_amount,
+        buy_amount: order_buy_amount,
+        network_fee: 20.ether().into_wei(),
     };
     let network_fee = 10.ether().into_wei();
-    let test_case = TestCase {
+    let test_case = PriceImprovementTestCase {
         order_side: order::Side::Sell,
-        fee_policy,
+        policy_factor: 0.5,
+        policy_max_volume_factor: 1.0,
+        quote,
         order_sell_amount,
+        order_buy_amount,
         network_fee: Some(network_fee),
-        quote_sell_amount: order_sell_amount,
-        quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
-        // order buy amount - quote network fee * factor (in sell token)
-        executed_buy_amount: 32.ether().into_wei(),
-        surplus_factor: 1.into(),
+        // todo: the value is changed after merge
+        executed_buy_amount: "34.285714285714285714".ether().into_wei(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    price_improvement_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_buy_in_market_order() {
-    let fee_policy = FeePolicy::PriceImprovement {
-        factor: 0.5,
-        max_volume_factor: 1.0,
-        quote: PriceImprovementQuote {
-            sell_amount: 50.ether().into_wei(),
-            buy_amount: 40.ether().into_wei(),
-            network_fee: 20.ether().into_wei(),
-        },
-    };
     let order_sell_amount = 100.ether().into_wei();
     let order_buy_amount = 40.ether().into_wei();
+    let quote = PriceImprovementQuote {
+        sell_amount: 50.ether().into_wei(),
+        buy_amount: 40.ether().into_wei(),
+        network_fee: 20.ether().into_wei(),
+    };
     let network_fee = 10.ether().into_wei();
-    let test_case = TestCase {
+    let test_case = PriceImprovementTestCase {
         order_side: order::Side::Buy,
-        fee_policy,
+        policy_factor: 0.5,
+        policy_max_volume_factor: 1.0,
+        quote,
         order_sell_amount,
+        order_buy_amount,
         network_fee: Some(network_fee),
-        quote_sell_amount: order_sell_amount,
-        quote_buy_amount: order_buy_amount,
         executed: order_buy_amount,
-        // todo: no price improvement since quote provides better conditions, but the fee isn't
-        // considered
+        // no price improvement since quote provides better conditions
         executed_sell_amount: order_sell_amount,
         executed_buy_amount: order_buy_amount,
-        surplus_factor: 1.into(),
     };
 
-    protocol_fee_test_case(test_case).await;
+    price_improvement_fee_test_case(test_case).await;
 }
 
 #[tokio::test]
 #[ignore]
 async fn price_improvement_fee_sell_in_market_order() {
-    let fee_policy = FeePolicy::PriceImprovement {
-        factor: 0.5,
-        max_volume_factor: 1.0,
-        quote: PriceImprovementQuote {
-            sell_amount: 50.ether().into_wei(),
-            buy_amount: 40.ether().into_wei(),
-            network_fee: 20.ether().into_wei(),
-        },
-    };
     let order_sell_amount: eth::U256 = 50.ether().into_wei();
     let order_buy_amount: eth::U256 = 10.ether().into_wei();
+    let quote = PriceImprovementQuote {
+        sell_amount: 50.ether().into_wei(),
+        buy_amount: 40.ether().into_wei(),
+        network_fee: 20.ether().into_wei(),
+    };
     let network_fee = 10.ether().into_wei();
-    let test_case = TestCase {
+    let test_case = PriceImprovementTestCase {
         order_side: order::Side::Sell,
-        fee_policy,
+        policy_factor: 0.5,
+        policy_max_volume_factor: 1.0,
+        quote,
         order_sell_amount,
+        order_buy_amount,
         network_fee: Some(network_fee),
-        quote_sell_amount: order_sell_amount,
-        quote_buy_amount: order_buy_amount,
         executed: order_sell_amount - network_fee,
-        // todo: no price improvement since quote provides better conditions, but the fee isn't
-        // considered
-        executed_sell_amount: 50.ether().into_wei(),
-        executed_buy_amount: 10.ether().into_wei(),
-        surplus_factor: 1.into(),
+        // no price improvement since quote provides better conditions
+        executed_sell_amount: order_sell_amount,
+        executed_buy_amount: order_buy_amount,
     };
 
-    protocol_fee_test_case(test_case).await;
+    price_improvement_fee_test_case(test_case).await;
 }

From 886fc2f80f4ccfc45200719c5001c8263d6941c7 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Fri, 8 Mar 2024 18:23:10 +0000
Subject: [PATCH 13/13] Wrong comment

---
 crates/driver/src/tests/cases/protocol_fees.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index a61fb13382..4df8aaa6ea 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -300,7 +300,7 @@ async fn price_improvement_fee_sell_out_of_market_order() {
         network_fee: Some(network_fee),
         executed: order_sell_amount - network_fee,
         executed_sell_amount: order_sell_amount,
-        // todo: the value is changed after merge
+        // todo: how to prove the value?
         executed_buy_amount: "34.285714285714285714".ether().into_wei(),
     };