feat: added option to load multiple api keys selected at random order

Author: Aviksaikat

.github/workflows/test.yaml

@@ -82,6 +82,7 @@ jobs:
           run: pytest -m "not fuzzing" -s --cov=src -n auto
           env:
             WEB3_INFURA_PROJECT_ID: ${{ secrets.WEB3_INFURA_PROJECT_ID }}
+            WEB3_INFURA_PROJECT_IDS: ${{ secrets.WEB3_INFURA_PROJECT_IDS }}
 
 # NOTE: uncomment this block after you've marked tests with @pytest.mark.fuzzing
 #    fuzzing:

ape_infura/provider.py

@@ -1,4 +1,5 @@
 import os
+import random
 from typing import Optional
 
 from ape.api import UpstreamProvider
@@ -9,7 +10,12 @@
 from web3.gas_strategies.rpc import rpc_gas_price_strategy
 from web3.middleware import geth_poa_middleware
 
-_ENVIRONMENT_VARIABLE_NAMES = ("WEB3_INFURA_PROJECT_ID", "WEB3_INFURA_API_KEY")
+_ENVIRONMENT_VARIABLE_NAMES = (
+    "WEB3_INFURA_PROJECT_ID",
+    "WEB3_INFURA_API_KEY",
+    "WEB3_INFURA_PROJECT_IDS",
+    "WEB3_INFURA_API_KEYS",
+)
 # NOTE: https://docs.infura.io/learn/websockets#supported-networks
 _WEBSOCKET_CAPABLE_ECOSYSTEMS = {
     "ethereum",
@@ -34,6 +40,27 @@ def __init__(self):
 
 class Infura(Web3Provider, UpstreamProvider):
     network_uris: dict[tuple[str, str], str] = {}
+    api_keys: list[str] = []
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.load_api_keys()
+
+    def load_api_keys(self):
+        self.api_keys = []
+        for env_var_name in _ENVIRONMENT_VARIABLE_NAMES:
+            env_var = os.environ.get(env_var_name)
+            if env_var:
+                if env_var_name.endswith("S"):  # Handle array-like environment variables
+                    self.api_keys.extend([key.strip() for key in env_var.split(",")])
+                else:
+                    self.api_keys.append(env_var)
+
+        if not self.api_keys:
+            raise MissingProjectKeyError()
+
+    def get_random_api_key(self):
+        return random.choice(self.api_keys)
 
     @property
     def uri(self) -> str:
@@ -42,15 +69,7 @@ def uri(self) -> str:
         if (ecosystem_name, network_name) in self.network_uris:
             return self.network_uris[(ecosystem_name, network_name)]
 
-        key = None
-        for env_var_name in _ENVIRONMENT_VARIABLE_NAMES:
-            env_var = os.environ.get(env_var_name)
-            if env_var:
-                key = env_var
-                break
-
-        if not key:
-            raise MissingProjectKeyError()
+        key = self.get_random_api_key()
 
         prefix = f"{ecosystem_name}-" if ecosystem_name != "ethereum" else ""
         network_uri = f"https://{prefix}{network_name}.infura.io/v3/{key}"

tests/test_provider.py

@@ -1,3 +1,5 @@
+from unittest.mock import patch
+
 import pytest
 import websocket  # type: ignore
 from ape.utils import ZERO_ADDRESS
@@ -31,3 +33,56 @@ def test_infura_ws(provider):
 
     except Exception as err:
         pytest.fail(f"Websocket URI not accessible. Reason: {err}")
+
+
+def test_load_multiple_api_keys(provider):
+    with patch.dict(
+        os.environ,
+        {"WEB3_INFURA_PROJECT_IDS": "key1,key2,key3", "WEB3_INFURA_API_KEYS": "key4,key5,key6"},
+    ):
+        provider.load_api_keys()
+        assert len(provider.api_keys) == 6
+        assert "key1" in provider.api_keys
+        assert "key6" in provider.api_keys
+
+
+def test_load_single_and_multiple_api_keys(provider):
+    with patch.dict(
+        os.environ,
+        {
+            "WEB3_INFURA_PROJECT_ID": "single_key1",
+            "WEB3_INFURA_API_KEY": "single_key2",
+            "WEB3_INFURA_PROJECT_IDS": "multi_key1,multi_key2",
+        },
+    ):
+        provider.load_api_keys()
+        assert len(provider.api_keys) == 4
+        assert "single_key1" in provider.api_keys
+        assert "multi_key2" in provider.api_keys
+
+
+def test_random_api_key_selection(provider):
+    with patch.dict(os.environ, {"WEB3_INFURA_PROJECT_IDS": "key1,key2,key3,key4,key5"}):
+        provider.load_api_keys()
+        selected_keys = set()
+        for _ in range(50):  # Run multiple times to ensure randomness
+            selected_keys.add(provider.get_random_api_key())
+        assert len(selected_keys) > 1  # Ensure we're getting different keys
+
+
+def test_uri_with_random_api_key(provider):
+    with patch.dict(os.environ, {"WEB3_INFURA_PROJECT_IDS": "key1,key2,key3"}):
+        provider.load_api_keys()
+        uris = set()
+        for _ in range(10):  # Generate multiple URIs
+            uri = provider.uri
+            uris.add(uri)
+            assert uri.startswith("https")
+            assert "/v3/key" in uri
+        assert len(uris) > 1  # Ensure we're getting different URIs with different keys
+
+
+def test_missing_project_key_error_raised(provider):
+    with patch.dict("os.environ", {}, clear=True):
+        with pytest.raises(MissingProjectKeyError):
+            provider.load_api_keys()