From cb6e9973755331eba08a1a73584b662eafc4921e Mon Sep 17 00:00:00 2001
From: Thales Lima <thales@datarevenue.com>
Date: Tue, 6 Aug 2024 17:28:47 +0200
Subject: [PATCH] Add token_factory

---
 testing/src/runner/runner.py | 15 ++++++--------
 testing/src/runner/utils.py  | 39 ++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/testing/src/runner/runner.py b/testing/src/runner/runner.py
index 40f7688..7324927 100644
--- a/testing/src/runner/runner.py
+++ b/testing/src/runner/runner.py
@@ -29,7 +29,7 @@ from tycho_client.stream import TychoStream
 from .adapter_handler import AdapterContractHandler
 from .evm import get_token_balance, get_block_header
 from .tycho import TychoRunner
-from .utils import build_snapshot_message
+from .utils import build_snapshot_message, token_factory
 
 
 class TestResult:
@@ -74,6 +74,7 @@ class TestRunner:
             db_url, with_binary_logs, self.config["initialized_accounts"]
         )
         self.tycho_rpc_client = TychoRPCClient()
+        self._token_factory_func = token_factory(self.tycho_rpc_client)
         self.db_url = db_url
         self._vm_traces = vm_traces
         self._chain = Chain.ethereum
@@ -244,14 +245,10 @@ class TestRunner:
                 )
 
             decoder = ThirdPartyPoolTychoDecoder(
-                adapter_contract=adapter_contract, minimum_gas=0, trace=self._vm_traces
-            )
-
-            stream_adapter = TychoStream(
-                tycho_url="0.0.0.0:4242",
-                exchanges=[protocol],
-                min_tvl=Decimal("0"),
-                blockchain=self._chain,
+                token_factory_func=self._token_factory_func,
+                adapter_contract=adapter_contract,
+                minimum_gas=0,
+                trace=self._vm_traces,
             )
 
             snapshot_message: Snapshot = build_snapshot_message(
diff --git a/testing/src/runner/utils.py b/testing/src/runner/utils.py
index 78cc9fc..eb34665 100644
--- a/testing/src/runner/utils.py
+++ b/testing/src/runner/utils.py
@@ -1,13 +1,20 @@
 from logging import getLogger
+from typing import Union
 
 from protosim_py.evm.pool_state import ThirdPartyPool
+from protosim_py.models import EthereumToken
 from tycho_client.dto import (
     ResponseProtocolState,
     ProtocolComponent,
     ResponseAccount,
     ComponentWithState,
     Snapshot,
+    HexBytes,
+    TokensParams,
+    PaginationParams,
+    ResponseToken,
 )
+from tycho_client.rpc_client import TychoRPCClient
 
 log = getLogger(__name__)
 
@@ -32,3 +39,35 @@ def build_snapshot_message(
 
     states = {id_: ComponentWithState(**state) for id_, state in states.items()}
     return Snapshot(states=states, vm_storage=vm_storage)
+
+
+def token_factory(rpc_client: TychoRPCClient) -> callable(HexBytes):
+    _client = rpc_client
+    _token_cache: dict[HexBytes, EthereumToken] = {}
+
+    def factory(addresses: Union[HexBytes, list[HexBytes]]) -> list[EthereumToken]:
+        if not isinstance(addresses, list):
+            addresses = [addresses]
+
+        response = dict()
+        to_fetch = []
+
+        for address in addresses:
+            if address in _token_cache:
+                response[address] = _token_cache[address]
+            else:
+                to_fetch.append(address)
+
+        if to_fetch:
+            pagination = PaginationParams(page_size=len(to_fetch), page=0)
+            params = TokensParams(token_addresses=to_fetch, pagination=pagination)
+            tokens = _client.get_tokens(params)
+            for token in tokens:
+                eth_token = EthereumToken(**token.dict())
+
+                response[token.address] = eth_token
+                _token_cache[token.address] = eth_token
+
+        return [response[address] for address in addresses]
+
+    return factory