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 =
+        let (balance_storage_addr, balance_slot) =
-            if let Some(Ok((_storage_addr, slot))) = results.get(&solution.given_token.clone()) {
+            if let Some(Ok((storage_addr, slot))) = results.get(&solution.given_token.clone()) {
-                slot
+                (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(
+        // Use the exact given amount for balance and allowance (no buffer, no max)
-            token_address,
+        let token_balance = biguint_to_u256(&solution.given_amount);
-            AccountOverride::default().with_state_diff(vec![
+        let token_allowance = biguint_to_u256(&solution.given_amount);
-                (
-                    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));
 
+        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!(
+                .wrap_err(format!("No state found for component: {id}"))?
-                    "Component {id} does not exist in protocol_states_by_id {protocol_states_by_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",

substreams/ethereum-balancer-v2/integration_test.tycho.yaml

@@ -1,6 +1,5 @@
 substreams_yaml_path: ./substreams.yaml
 protocol_system: "vm:balancer_v2"
-module_name: "map_protocol_changes"
 protocol_type_names:
   - "balancer_v2_pool"
 adapter_contract: "BalancerV2SwapAdapter"

substreams/ethereum-curve/integration_test.tycho.yaml

@@ -1,6 +1,5 @@
 substreams_yaml_path: ./substreams.yaml
 protocol_system: "vm:curve"
-module_name: "map_protocol_changes"
 protocol_type_names:
   - "curve_pool"
 adapter_contract: "CurveAdapter"

substreams/ethereum-maverick-v2/integration_test.tycho.yaml

@@ -1,6 +1,5 @@
 substreams_yaml_path: ./ethereum-maverick-v2.yaml
 protocol_system: "vm:maverick_v2"
-module_name: "map_protocol_changes"
 adapter_contract: "MaverickV2SwapAdapter"
 adapter_build_signature: "constructor(address,address)"
 adapter_build_args: "0x0A7e848Aca42d879EF06507Fca0E7b33A0a63c1e,0xb40AfdB85a07f37aE217E7D6462e609900dD8D7A"

substreams/ethereum-pancakeswap-v3/integration_test.tycho.yaml

@@ -1,5 +1,4 @@
 substreams_yaml_path: ./ethereum-pancakeswap-v3.yaml
-protocol_system: "pancakeswap_v3"
 protocol_type_names:
   - "pancakeswap_v3_pool"
 module_name: "map_protocol_changes"

substreams/ethereum-template-factory/integration_test.tycho.yaml

@@ -6,7 +6,7 @@ adapter_contract: "SwapAdapter"
 adapter_build_signature: "constructor(address)"
 # A comma separated list of args to be passed to the contructor of the Adapter contract"
 adapter_build_args: "0x0000000000000000000000000000000000000000"
-# Whether or not the testing script should skip checking balances of the protocol components.
+# Whether the testing script should skip checking balances of the protocol components.
 # If set to `true` please always add a reason why it's skipped.
 skip_balance_check: false
 # Accounts that will be automatically initialized at test start
@@ -29,6 +29,8 @@ protocol_type_names:
   - "type_name_1"
   - "type_name_2"
 # A list of tests.
+# The name of the protocol system
+protocol_system: "protocol_name"
 tests:
   # Name of the test
   - name: test_pool_creation
@@ -45,11 +47,14 @@ tests:
           - "0xdac17f958d2ee523a2206206994597c13d831ec7"
           - "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48"
           - "0x6b175474e89094c44da98b954eedeac495271d0f"
-        static_attributes: {}
+        static_attributes: { }
         creation_tx: "0x20793bbf260912aae189d5d261ff003c9b9166da8191d8f9d63ff1c7722f3ac6"
-        # Whether or not the script should skip trying to simulate a swap on this component.
+        # Whether the script should skip trying to simulate a swap on this component.
         # If set to `true` please always add a reason why it's skipped.
         skip_simulation: false
+        # Whether the script should skip trying to simulate execution of a swap on this component.
+        # If set to `true` please always add a reason why it's skipped.
+        skip_execution: false
   - name: test_something_else
     start_block: 123
     stop_block: 456
@@ -58,6 +63,7 @@ tests:
         tokens:
           - "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE"
           - "0xae7ab96520DE3A18E5e111B5EaAb095312D7fE84"
-        static_attributes: {}
+        static_attributes: { }
         creation_tx: "0xfac67ecbd423a5b915deff06045ec9343568edaec34ae95c43d35f2c018afdaa"
         skip_simulation: true # If true, always add a reason
+        skip_execution: true # If true, always add a reason

substreams/ethereum-template-singleton/integration_test.tycho.yaml

@@ -6,7 +6,7 @@ adapter_contract: "SwapAdapter"
 adapter_build_signature: "constructor(address)"
 # A comma separated list of args to be passed to the contructor of the Adapter contract"
 adapter_build_args: "0x0000000000000000000000000000000000000000"
-# Whether or not the testing script should skip checking balances of the protocol components.
+# Whether the testing script should skip checking balances of the protocol components.
 # If set to `true` please always add a reason why it's skipped.
 skip_balance_check: false
 # A list of accounts that need to be indexed to run the tests properly.
@@ -20,6 +20,8 @@ initialized_accounts:
 protocol_type_names:
   - "type_name_1"
   - "type_name_2"
+# The name of the protocol system
+protocol_system: "protocol_name"
 # A list of tests.
 tests:
   # Name of the test
@@ -37,11 +39,14 @@ tests:
           - "0xdac17f958d2ee523a2206206994597c13d831ec7"
           - "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48"
           - "0x6b175474e89094c44da98b954eedeac495271d0f"
-        static_attributes: {}
+        static_attributes: { }
         creation_tx: "0x20793bbf260912aae189d5d261ff003c9b9166da8191d8f9d63ff1c7722f3ac6"
-        # Whether or not the script should skip trying to simulate a swap on this component.
+        # Whether the script should skip trying to simulate a swap on this component.
         # If set to `true` please always add a reason why it's skipped.
         skip_simulation: false
+        # Whether the script should skip trying to simulate execution of a swap on this component.
+        # If set to `true` please always add a reason why it's skipped.
+        skip_execution: false
   - name: test_something_else
     start_block: 123
     stop_block: 456
@@ -50,6 +55,7 @@ tests:
         tokens:
           - "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE"
           - "0xae7ab96520DE3A18E5e111B5EaAb095312D7fE84"
-        static_attributes: {}
+        static_attributes: { }
         creation_tx: "0xfac67ecbd423a5b915deff06045ec9343568edaec34ae95c43d35f2c018afdaa"
         skip_simulation: true # If true, always add a reason
+        skip_execution: true # If true, always add a reason