feat: support env var replacements (#16)

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"]

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