From 07315294ce1d058c3f0231f13ccd8b119ccf927d Mon Sep 17 00:00:00 2001
From: zach <zachchou016@gmail.com>
Date: Tue, 7 Oct 2025 17:56:55 +0200
Subject: [PATCH] fix: add traced entrypoint

---
 protocol-testing/src/test_runner.rs | 29 +++++++++++++++++++++++++++--
 protocol-testing/src/tycho_rpc.rs   | 27 +++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/protocol-testing/src/test_runner.rs b/protocol-testing/src/test_runner.rs
index 87717da..51c96aa 100644
--- a/protocol-testing/src/test_runner.rs
+++ b/protocol-testing/src/test_runner.rs
@@ -30,7 +30,10 @@ use tycho_simulation::{
         BlockHeader, FeedMessage,
     },
     tycho_common::{
-        dto::{Chain, ProtocolComponent, ResponseAccount, ResponseProtocolState},
+        dto::{
+            Chain, EntryPointWithTracingParams, ProtocolComponent, ResponseAccount,
+            ResponseProtocolState, TracingResult,
+        },
         models::token::Token,
         Bytes,
     },
@@ -369,6 +372,16 @@ impl TestRunner {
             .into_diagnostic()
             .wrap_err("Failed to get contract state")?;
 
+        let traced_entry_points = self
+            .runtime
+            .block_on(tycho_client.get_traced_entry_points(
+                protocol_system,
+                expected_component_ids.clone(),
+                chain,
+            ))
+            .into_diagnostic()
+            .wrap_err("Failed to get trace points")?;
+
         // Create a map of component IDs to components for easy lookup
         let mut components_by_id: HashMap<String, ProtocolComponent> = protocol_components
             .clone()
@@ -386,6 +399,7 @@ impl TestRunner {
 
         debug!("Found {} protocol components", components_by_id.len());
         debug!("Found {} protocol states", protocol_states_by_id.len());
+        debug!("Found {} traced entry points", traced_entry_points.len());
 
         let adapter_contract_path;
         let mut adapter_contract_path_str: Option<&str> = None;
@@ -441,13 +455,24 @@ impl TestRunner {
                 .wrap_err(format!("No state found for component: {id}"))?
                 .clone();
 
+            let traced_entry_points: Vec<(EntryPointWithTracingParams, TracingResult)> =
+                traced_entry_points
+                    .get(component_id)
+                    .map(|inner| {
+                        inner
+                            .iter()
+                            .map(|(k, v)| (k.clone(), v.clone()))
+                            .collect::<Vec<_>>()
+                    })
+                    .unwrap_or_default();
+
             let component_with_state = ComponentWithState {
                 state,
                 component: component.clone(),
                 component_tvl: None,
                 // Neither UniswapV4 with hooks not certain balancer pools are currently supported
                 // for SDK testing
-                entrypoints: vec![],
+                entrypoints: traced_entry_points,
             };
             states.insert(component_id.clone(), component_with_state);
         }
diff --git a/protocol-testing/src/tycho_rpc.rs b/protocol-testing/src/tycho_rpc.rs
index a1cce71..c268501 100644
--- a/protocol-testing/src/tycho_rpc.rs
+++ b/protocol-testing/src/tycho_rpc.rs
@@ -5,8 +5,9 @@ use tycho_simulation::{
     tycho_client::{rpc::RPCClient, HttpRPCClient},
     tycho_common::{
         dto::{
-            Chain, PaginationParams, ProtocolComponent, ProtocolComponentsRequestBody,
-            ResponseAccount, ResponseProtocolState, ResponseToken, StateRequestBody, VersionParam,
+            Chain, EntryPointWithTracingParams, PaginationParams, ProtocolComponent,
+            ProtocolComponentsRequestBody, ResponseAccount, ResponseProtocolState, ResponseToken,
+            StateRequestBody, TracedEntryPointRequestBody, TracingResult, VersionParam,
         },
         models::token::Token,
         Bytes,
@@ -153,4 +154,26 @@ impl TychoClient {
 
         Ok(res)
     }
+
+    /// Gets traced entry points from the RPC server
+    pub async fn get_traced_entry_points(
+        &self,
+        protocol_system: &str,
+        component_ids: Vec<String>,
+        chain: Chain,
+    ) -> Result<HashMap<String, Vec<(EntryPointWithTracingParams, TracingResult)>>, RpcError> {
+        let request_body = TracedEntryPointRequestBody {
+            protocol_system: protocol_system.to_string(),
+            chain,
+            pagination: PaginationParams { page: 0, page_size: 100 },
+            component_ids: Some(component_ids),
+        };
+
+        let traced_entry_points = self
+            .http_client
+            .get_traced_entry_points(&request_body)
+            .await?;
+
+        Ok(traced_entry_points.traced_entry_points)
+    }
 }