fix: Misc improvements (#291)

* fix: Misc improvements

Make module_name optional and default to map_protocol_changes
Add build instructions for the DB image inside of the docker-compose
Improve error message
Remove unnecessary module_name from yaml files

#time 39m

* fix: Bring back module_name in balancer v3 tests

#time 8m

* fix: Fix balance and allowance overwrites

For tokens with proxies we need to use the address that is returned by the detector

#time 13m

protocol-testing/docker-compose.yaml

@@ -1,5 +1,8 @@
 services:
   db:
+    build:
+      context: .
+      dockerfile: postgres.Dockerfile
     image: protocol-testing-db:latest
     restart: "always"
     healthcheck:
@@ -30,7 +33,7 @@ services:
       SUBSTREAMS_API_TOKEN: "${SUBSTREAMS_API_TOKEN}"
     # PROTOCOLS to test separated by space and with optional filter
     # e.g. "ethereum-balancer-v2=weighted_legacy_creation ethereum-ekubo-v2"
-    entrypoint: ["/entrypoint.sh", "${PROTOCOLS}"]
+    entrypoint: [ "/entrypoint.sh", "${PROTOCOLS}" ]
 
 volumes:
   postgres_data:

protocol-testing/src/config.rs

@@ -149,6 +149,6 @@ pub struct IntegrationTestsConfig {
     pub skip_balance_check: bool,
     pub protocol_type_names: Vec<String>,
     pub protocol_system: String,
-    pub module_name: String,
+    pub module_name: Option<String>,
     pub tests: Vec<IntegrationTest>,
 }

protocol-testing/src/execution.rs

@@ -242,13 +242,12 @@ async fn setup_user_overwrites(
     block: &Block,
 ) -> miette::Result<AddressHashMap<AccountOverride>> {
     let mut overwrites = AddressHashMap::default();
-    // Add ETH balance override for the user to ensure they have enough gas funds
-    let mut eth_balance = U256::from_str("100000000000000000000").unwrap(); // 100 ETH
-
     let token_address = Address::from_slice(&solution.given_token[..20]);
-    // If given token is ETH, add the given amount to the balance
+    // If given token is ETH, add the given amount to the balance + 100 ETH for gas
     if solution.given_token == Bytes::zero(20) {
-        eth_balance += biguint_to_u256(&solution.given_amount);
+        let eth_balance = biguint_to_u256(&solution.given_amount) +
+            U256::from_str("100000000000000000000").unwrap(); // given_amount + 100 ETH for gas
+        overwrites.insert(user_address, AccountOverride::default().with_balance(eth_balance));
     // if the given token is not ETH, do balance and allowance slots overwrites
     } else {
         let detector = EVMBalanceSlotDetector::new(BalanceSlotDetectorConfig {
@@ -265,9 +264,9 @@ async fn setup_user_overwrites(
             )
             .await;
 
-        let balance_slot =
-            if let Some(Ok((_storage_addr, slot))) = results.get(&solution.given_token.clone()) {
-                slot
+        let (balance_storage_addr, balance_slot) =
+            if let Some(Ok((storage_addr, slot))) = results.get(&solution.given_token.clone()) {
+                (storage_addr, slot)
             } else {
                 return Err(miette!("Couldn't find balance storage slot for token {token_address}"));
             };
@@ -287,32 +286,59 @@ async fn setup_user_overwrites(
             )
             .await;
 
-        let allowance_slot = if let Some(Ok((_storage_addr, slot))) =
+        let (allowance_storage_addr, allowance_slot) = if let Some(Ok((storage_addr, slot))) =
             results.get(&solution.given_token.clone())
         {
-            slot
+            (storage_addr, slot)
         } else {
             return Err(miette!("Couldn't find allowance storage slot for token {token_address}"));
         };
 
-        overwrites.insert(
-            token_address,
-            AccountOverride::default().with_state_diff(vec![
-                (
-                    alloy::primitives::B256::from_slice(allowance_slot),
-                    alloy::primitives::B256::from_slice(&U256::MAX.to_be_bytes::<32>()),
-                ),
-                (
-                    alloy::primitives::B256::from_slice(balance_slot),
-                    alloy::primitives::B256::from_slice(
-                        &biguint_to_u256(&solution.given_amount).to_be_bytes::<32>(),
-                    ),
-                ),
-            ]),
-        );
-    }
-    overwrites.insert(user_address, AccountOverride::default().with_balance(eth_balance));
+        // Use the exact given amount for balance and allowance (no buffer, no max)
+        let token_balance = biguint_to_u256(&solution.given_amount);
+        let token_allowance = biguint_to_u256(&solution.given_amount);
 
+        let balance_storage_address = Address::from_slice(&balance_storage_addr[..20]);
+        let allowance_storage_address = Address::from_slice(&allowance_storage_addr[..20]);
+
+        // Apply balance and allowance overrides
+        // If both storage addresses are the same, combine them into one override
+        if balance_storage_address == allowance_storage_address {
+            overwrites.insert(
+                balance_storage_address,
+                AccountOverride::default().with_state_diff(vec![
+                    (
+                        alloy::primitives::B256::from_slice(balance_slot),
+                        alloy::primitives::B256::from_slice(&token_balance.to_be_bytes::<32>()),
+                    ),
+                    (
+                        alloy::primitives::B256::from_slice(allowance_slot),
+                        alloy::primitives::B256::from_slice(&token_allowance.to_be_bytes::<32>()),
+                    ),
+                ]),
+            );
+        } else {
+            // Different storage addresses, apply separately
+            overwrites.insert(
+                balance_storage_address,
+                AccountOverride::default().with_state_diff(vec![(
+                    alloy::primitives::B256::from_slice(balance_slot),
+                    alloy::primitives::B256::from_slice(&token_balance.to_be_bytes::<32>()),
+                )]),
+            );
+            overwrites.insert(
+                allowance_storage_address,
+                AccountOverride::default().with_state_diff(vec![(
+                    alloy::primitives::B256::from_slice(allowance_slot),
+                    alloy::primitives::B256::from_slice(&token_allowance.to_be_bytes::<32>()),
+                )]),
+            );
+        }
+
+        // Add 1 ETH for gas for non-ETH token swaps
+        let eth_balance = U256::from_str("100000000000000000000").unwrap(); // 100 ETH for gas
+        overwrites.insert(user_address, AccountOverride::default().with_balance(eth_balance));
+    }
     Ok(overwrites)
 }
 

protocol-testing/src/test_runner.rs

@@ -234,7 +234,7 @@ impl TestRunner {
                 test.stop_block,
                 &config.protocol_type_names,
                 &config.protocol_system,
-                &config.module_name,
+                config.module_name.clone(),
             )
             .wrap_err("Failed to run Tycho")?;
 
@@ -438,9 +438,7 @@ impl TestRunner {
 
             let state = protocol_states_by_id
                 .get(component_id)
-                .wrap_err(format!(
-                    "Component {id} does not exist in protocol_states_by_id {protocol_states_by_id:?}"
-                ))?
+                .wrap_err(format!("No state found for component: {id}"))?
                 .clone();
 
             let component_with_state = ComponentWithState {

protocol-testing/src/tycho_runner.rs

@@ -31,7 +31,7 @@ impl TychoRunner {
         end_block: u64,
         protocol_type_names: &[String],
         protocol_system: &str,
-        module_name: &str,
+        module_name: Option<String>,
     ) -> miette::Result<()> {
         info!("Running Tycho indexer from block {start_block} to {end_block}...");
 
@@ -48,7 +48,9 @@ impl TychoRunner {
             "--spkg",
             spkg_path,
             "--module",
-            module_name,
+            module_name
+                .as_deref()
+                .unwrap_or("map_protocol_changes"),
             "--protocol-type-names",
             &protocol_type_names.join(","),
             "--protocol-system",