From 2ac00e79749f67a2717207181a95bdb18d528873 Mon Sep 17 00:00:00 2001
From: ilya <ilya@cow.fi>
Date: Thu, 14 Mar 2024 18:47:23 +0000
Subject: [PATCH] Applied suggestion

---
 .../driver/src/tests/cases/protocol_fees.rs   | 47 +++++----------
 crates/driver/src/tests/setup/blockchain.rs   | 59 ++++++++++---------
 crates/driver/src/tests/setup/mod.rs          |  9 ++-
 3 files changed, 54 insertions(+), 61 deletions(-)

diff --git a/crates/driver/src/tests/cases/protocol_fees.rs b/crates/driver/src/tests/cases/protocol_fees.rs
index 8efb6544f7..4814f4ad11 100644
--- a/crates/driver/src/tests/cases/protocol_fees.rs
+++ b/crates/driver/src/tests/cases/protocol_fees.rs
@@ -10,7 +10,6 @@ use crate::{
             ab_solution,
             fee::{Policy, Quote},
             ExpectedOrderAmounts,
-            Partial,
             Test,
         },
     },
@@ -52,37 +51,23 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         .sell_amount(test_case.execution.solver.sell)
         .buy_amount(test_case.execution.solver.buy);
     let pool = ab_adjusted_pool(quote);
-    // Check if the order is expected to be partially filled by calculating a
-    // difference between solver's and order's target amounts
-    let partially_executed = match test_case.order.side {
-        order::Side::Sell => test_case
-            .order
-            .sell_amount
-            .saturating_sub(test_case.execution.solver.sell),
-        order::Side::Buy => test_case
-            .order
-            .buy_amount
-            .saturating_sub(test_case.execution.solver.buy),
-    };
-    // If there is a difference, the order is considered to be partially fillable
-    let partial = if partially_executed <= eth::U256::zero() {
-        Partial::No
-    } else {
-        Partial::Yes {
-            executed: eth::U256::zero(),
-        }
-    };
     let solver_fee = test_case.execution.driver.sell / 100;
-    // Target amount to be executed by the solver in case of partially fillable
-    // order
-    let executed = match partial {
-        Partial::Yes { .. } => match test_case.order.side {
-            order::Side::Buy => Some(test_case.execution.solver.buy),
-            order::Side::Sell => Some(test_case.execution.solver.sell - solver_fee),
-        },
-        Partial::No => None,
+    let executed = match test_case.order.side {
+        order::Side::Buy => {
+            if test_case.order.buy_amount > test_case.execution.solver.buy {
+                Some(test_case.execution.solver.buy)
+            } else {
+                None
+            }
+        }
+        order::Side::Sell => {
+            if test_case.order.sell_amount > test_case.execution.solver.sell {
+                Some(test_case.execution.solver.sell - solver_fee)
+            } else {
+                None
+            }
+        }
     };
-
     // Amounts expected to be returned by the driver after fee processing
     let expected_amounts = ExpectedOrderAmounts {
         sell: test_case.execution.driver.sell,
@@ -99,7 +84,7 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         .side(test_case.order.side)
         .fee_policy(test_case.fee_policy)
         .executed(executed)
-        .partial(partial)
+        .partial(0.into())
         // Surplus is configured explicitly via executed/quoted amounts
         .no_surplus()
         .expected_amounts(expected_amounts);
diff --git a/crates/driver/src/tests/setup/blockchain.rs b/crates/driver/src/tests/setup/blockchain.rs
index dab8b2b9d2..1fa49d7f60 100644
--- a/crates/driver/src/tests/setup/blockchain.rs
+++ b/crates/driver/src/tests/setup/blockchain.rs
@@ -561,39 +561,42 @@ impl Blockchain {
     /// Compute the execution of an order given the available liquidity
     pub fn execution(&self, order: &Order) -> Execution {
         let pair = self.find_pair(order);
-        let (sell, buy) = match (order.side, order.executed.or(order.buy_amount)) {
-            // For buy order with explicitly specified amounts, use the buy amount
-            (order::Side::Buy, Some(buy_amount)) => (
-                pair.pool.in_given_out(Asset {
-                    amount: buy_amount,
-                    token: order.buy_token,
-                }),
-                buy_amount,
-            ),
-            // todo: simplify it
-            (order::Side::Sell, _) => {
-                let sell_amount = order
+        match order.side {
+            order::Side::Buy => {
+                // For buy order with explicitly specified amounts, use the buy amount,
+                // otherwise assume the full sell amount to compute the execution
+                let executed = order.executed.or(order.buy_amount);
+                match executed {
+                    Some(executed) => Execution {
+                        buy: executed,
+                        sell: pair.pool.in_given_out(Asset {
+                            amount: executed,
+                            token: order.buy_token,
+                        }),
+                    },
+                    None => Execution {
+                        buy: pair.pool.out_given_in(Asset {
+                            amount: order.sell_amount,
+                            token: order.sell_token,
+                        }),
+                        sell: order.sell_amount,
+                    },
+                }
+            }
+            order::Side::Sell => {
+                let executed = order
                     .executed
-                    .map(|sell| sell + order.solver_fee.unwrap_or_default())
+                    .map(|amount| amount + order.solver_fee.unwrap_or_default())
                     .unwrap_or(order.sell_amount);
-                (
-                    sell_amount,
-                    pair.pool.out_given_in(Asset {
-                        amount: sell_amount,
+                Execution {
+                    buy: pair.pool.out_given_in(Asset {
+                        amount: executed,
                         token: order.sell_token,
                     }),
-                )
+                    sell: executed,
+                }
             }
-            // Otherwise assume the full sell amount to compute the execution
-            (_, _) => (
-                order.sell_amount,
-                pair.pool.out_given_in(Asset {
-                    amount: order.sell_amount,
-                    token: order.sell_token,
-                }),
-            ),
-        };
-        Execution { sell, buy }
+        }
     }
 
     /// Set up the blockchain context and return the interactions needed to
diff --git a/crates/driver/src/tests/setup/mod.rs b/crates/driver/src/tests/setup/mod.rs
index abb0787472..5bbb6c9a0a 100644
--- a/crates/driver/src/tests/setup/mod.rs
+++ b/crates/driver/src/tests/setup/mod.rs
@@ -254,8 +254,13 @@ impl Order {
         }
     }
 
-    pub fn partial(self, partial: Partial) -> Self {
-        Self { partial, ..self }
+    pub fn partial(self, already_executed: eth::U256) -> Self {
+        Self {
+            partial: Partial::Yes {
+                executed: already_executed,
+            },
+            ..self
+        }
     }
 
     pub fn executed(self, executed: Option<eth::U256>) -> Self {