Load Snapshot from RPC request

This commit is contained in:
Thales Lima
2024-07-18 04:54:02 +02:00
committed by tvinagre
parent 5e6c7d4647
commit 183868e536
8 changed files with 500 additions and 59 deletions

View File

@@ -1,12 +1,23 @@
import itertools
import itertools
import os
from pathlib import Path
import shutil
import subprocess
from collections import defaultdict
from datetime import datetime
from decimal import Decimal
from pathlib import Path
import yaml
from pydantic import BaseModel
from evm import get_token_balance
from evm import get_token_balance, get_block_header
from tycho import TychoRunner
from tycho_client.tycho.decoders import ThirdPartyPoolTychoDecoder
from tycho_client.tycho.models import Blockchain, EVMBlock
from tycho_client.tycho.tycho_adapter import (
TychoPoolStateStreamAdapter,
)
class TestResult:
@@ -29,12 +40,20 @@ def load_config(yaml_path: str) -> dict:
return yaml.safe_load(file)
class SimulationFailure(BaseModel):
pool_id: str
sell_token: str
buy_token: str
error: str
class TestRunner:
def __init__(self, config_path: str, with_binary_logs: bool, db_url: str):
self.config = load_config(config_path)
self.base_dir = os.path.dirname(config_path)
self.tycho_runner = TychoRunner(with_binary_logs)
self.db_url = db_url
self._chain = Blockchain.ethereum
def run_tests(self) -> None:
"""Run all tests specified in the configuration."""
@@ -58,12 +77,11 @@ class TestRunner:
if result.success:
print(f"{test['name']} passed.")
else:
print(f"❗️ {test['name']} failed: {result.message}")
self.tycho_runner.empty_database(
self.db_url
)
self.tycho_runner.empty_database(self.db_url)
def validate_state(self, expected_state: dict, stop_block: int) -> TestResult:
"""Validate the current protocol state against the expected state."""
@@ -90,7 +108,7 @@ class TestRunner:
)
if isinstance(value, list):
if set(map(str.lower, value)) != set(
map(str.lower, component[key])
map(str.lower, component[key])
):
return TestResult.Failed(
f"List mismatch for key '{key}': {value} != {component[key]}"
@@ -100,25 +118,116 @@ class TestRunner:
f"Value mismatch for key '{key}': {value} != {component[key]}"
)
token_balances: dict[str, dict[str, int]] = defaultdict(dict)
for component in protocol_components["protocol_components"]:
comp_id = component["id"].lower()
for token in component["tokens"]:
token_lower = token.lower()
state = next((s for s in protocol_states["states"] if s["component_id"].lower() == comp_id), None)
state = next(
(
s
for s in protocol_states["states"]
if s["component_id"].lower() == comp_id
),
None,
)
if state:
balance_hex = state["balances"].get(token_lower, "0x0")
else:
balance_hex = "0x0"
tycho_balance = int(balance_hex, 16)
token_balances[comp_id][token_lower] = tycho_balance
node_balance = get_token_balance(token, comp_id, stop_block)
tycho_balance = int(balance_hex, 16)
if node_balance != tycho_balance:
return TestResult.Failed(
f"Balance mismatch for {comp_id}:{token} at block {stop_block}: got {node_balance} from rpc call and {tycho_balance} from Substreams")
f"Balance mismatch for {comp_id}:{token} at block {stop_block}: got {node_balance} from rpc call and {tycho_balance} from Substreams"
)
contract_states = self.tycho_runner.get_contract_state()
self.simulate_get_amount_out(
token_balances,
stop_block,
protocol_states,
protocol_components,
contract_states,
)
return TestResult.Passed()
except Exception as e:
return TestResult.Failed(str(e))
def simulate_get_amount_out(
self,
token_balances: dict[str, dict[str, int]],
block_number: int,
protocol_states: dict,
protocol_components: dict,
contract_state: dict,
) -> TestResult:
protocol_type_names = self.config["protocol_type_names"]
block_header = get_block_header(block_number)
block: EVMBlock = EVMBlock(
id=block_number,
ts=datetime.fromtimestamp(block_header.timestamp),
hash_=block_header.hash.hex(),
)
failed_simulations = dict[str, list[SimulationFailure]]
for protocol in protocol_type_names:
# TODO: Parametrize this
decoder = ThirdPartyPoolTychoDecoder(
"CurveSwapAdapter.evm.runtime", 0, False
)
stream_adapter = TychoPoolStateStreamAdapter(
tycho_url="0.0.0.0:4242",
protocol=protocol,
decoder=decoder,
blockchain=self._chain,
)
snapshot_message = stream_adapter.build_snapshot_message(
protocol_components, protocol_states, contract_state
)
decoded = stream_adapter.process_snapshot(block, snapshot_message)
for pool_state in decoded.pool_states.values():
pool_id = pool_state.id_
protocol_balances = token_balances.get(pool_id)
if not protocol_balances:
raise ValueError(f"Missing balances for pool {pool_id}")
for sell_token, buy_token in itertools.permutations(
pool_state.tokens, 2
):
try:
# Try to sell 0.1% of the protocol balance
sell_amount = Decimal("0.001") * sell_token.from_onchain_amount(
protocol_balances[sell_token.address]
)
amount_out, gas_used, _ = pool_state.get_amount_out(
sell_token, sell_amount, buy_token
)
# TODO: Should we validate this with an archive node or RPC reader?
print(
f"Amount out for {pool_id}: {sell_amount} {sell_token} -> {amount_out} {buy_token} - "
f"Gas used: {gas_used}"
)
except Exception as e:
print(
f"Error simulating get_amount_out for {pool_id}: {sell_token} -> {buy_token}. "
f"Error: {e}"
)
if pool_id not in failed_simulations:
failed_simulations[pool_id] = []
failed_simulations[pool_id].append(
SimulationFailure(
pool_id=pool_id,
sell_token=sell_token,
buy_token=buy_token,
error=str(e),
)
)
continue
@staticmethod
def build_spkg(yaml_file_path: str, modify_func: callable) -> str:
"""Build a Substreams package with modifications to the YAML file."""