diff --git a/evmchains/__init__.py b/evmchains/__init__.py index 51961f3..82b4dae 100644 --- a/evmchains/__init__.py +++ b/evmchains/__init__.py @@ -1,8 +1,13 @@ +import os import random +import re +from typing import List from evmchains.chains import PUBLIC_CHAIN_META from evmchains.types import Chain +ENV_VAR_REGEX = re.compile(r"\$\{([A-Za-z0-9_]+)\}") + def get_chain_meta(ecosystem: str, network: str) -> Chain: """Return a Chain instance with metadata for an EVM chain""" @@ -15,10 +20,31 @@ def get_chain_meta(ecosystem: str, network: str) -> Chain: return Chain(**PUBLIC_CHAIN_META[ecosystem][network]) +def get_rpcs(ecosystem: str, network: str) -> List[str]: + """Get a list of valid RPC endpoints for an ecosystem:network pair""" + rpcs = [] + + chain = get_chain_meta(ecosystem, network) + + for rpc in chain.rpc: + # Look for env var replacements in the endpoint URL + match = ENV_VAR_REGEX.search(rpc) + if match: + to_replace = match.group(0) + env_var = match.group(1) + # Include them only if the env var is available + if env_var in os.environ: + rpcs.append(rpc.replace(to_replace, os.environ[env_var])) + else: + rpcs.append(rpc) + + return rpcs + + def get_random_rpc(ecosystem: str, network: str) -> str: """Return a random RPC endpoint for an ecosystem:network pair""" - chain = get_chain_meta(ecosystem, network) - return random.choice(chain.rpc) + rpcs = get_rpcs(ecosystem, network) + return random.choice(rpcs) -__all__ = ["PUBLIC_CHAIN_META", "Chain", "get_chain_meta"] +__all__ = ["PUBLIC_CHAIN_META", "Chain", "get_chain_meta", "get_random_rpc", "get_rpcs"] diff --git a/tests/test_func.py b/tests/test_func.py index 1a537b5..ff41fbd 100644 --- a/tests/test_func.py +++ b/tests/test_func.py @@ -1,4 +1,26 @@ -from evmchains import PUBLIC_CHAIN_META, Chain, get_chain_meta +import os +from contextlib import contextmanager + +import pytest + +from evmchains import PUBLIC_CHAIN_META, Chain, get_chain_meta, get_rpcs + + +@pytest.fixture() +def env(): + @contextmanager + def inner(environ=dict(), unset=list()): + old_environ = os.environ.copy() + os.environ.update(environ) + for key in unset: + try: + del os.environ[key] + except KeyError: + pass + yield + os.environ = old_environ + + return inner def test_chain_structure(): @@ -26,3 +48,41 @@ def test_get_chain_meta(): assert chain.name == "Arbitrum One" assert chain.nativeCurrency["symbol"] == "ETH" assert len(chain.explorers) > 0 + + +def test_get_rpcs_alchemy(env): + alchemy_key = "alchemistsayswhat" + expected_rpc = f"https://arb-mainnet.g.alchemy.com/v2/{alchemy_key}" + + def find_rpc(rpcs, expected): + for rpc in rpcs: + if rpc == expected: + return True + return False + + with env({"ALCHEMY_API_KEY": alchemy_key}, unset=["INFURA_API_KEY"]): + rpcs = get_rpcs("arbitrum", "mainnet") + assert isinstance(rpcs, list) + assert len(rpcs) > 0 + assert find_rpc(rpcs, expected_rpc) + # Ensure the infura endpoint was also filtered out + assert len([rpc for rpc in rpcs if "infura" in rpc]) == 0 + + +def test_get_rpcs_infura(env): + infura_key = "infuraaaaaa" + expected_rpc = f"https://arbitrum-mainnet.infura.io/v3/{infura_key}" + + def find_rpc(rpcs, expected): + for rpc in rpcs: + if rpc == expected: + return True + return False + + with env({"INFURA_API_KEY": infura_key}, unset=["ALCHEMY_API_KEY"]): + rpcs = get_rpcs("arbitrum", "mainnet") + assert isinstance(rpcs, list) + assert len(rpcs) > 0 + assert find_rpc(rpcs, expected_rpc) + # Ensure the alchemy endpoint was also filtered out + assert len([rpc for rpc in rpcs if "alchemy" in rpc]) == 0