From 86e575fc7d08fe669c27f7c053966d17553af9df Mon Sep 17 00:00:00 2001
From: Saikat Karmakar <31238298+Aviksaikat@users.noreply.github.com>
Date: Mon, 19 Aug 2024 22:15:00 +0530
Subject: [PATCH] feat: added option to load multiple api keys selected at
 random order (#88)

---
 README.md              |  3 +++
 ape_infura/provider.py | 38 +++++++++++++++++++++++++++++---------
 setup.py               |  1 +
 tests/test_provider.py | 41 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 74 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 1d177c1..9434621 100644
--- a/README.md
+++ b/README.md
@@ -45,6 +45,9 @@ Either in your current terminal session or in your root RC file (e.g. `.bashrc`)
 
 ```bash
 export WEB3_INFURA_PROJECT_ID=MY_API_TOKEN
+
+# Multple tokens
+export WEB3_INFURA_PROJECT_ID=MY_API_TOKEN1, MY_API_TOKEN2
 ```
 
 To use the Infura provider plugin in most commands, set it via the `--network` option:
diff --git a/ape_infura/provider.py b/ape_infura/provider.py
index 415dfc1..37f681e 100644
--- a/ape_infura/provider.py
+++ b/ape_infura/provider.py
@@ -1,4 +1,5 @@
 import os
+import random
 from typing import Optional
 
 from ape.api import UpstreamProvider
@@ -34,6 +35,26 @@ def __init__(self):
 
 class Infura(Web3Provider, UpstreamProvider):
     network_uris: dict[tuple[str, str], str] = {}
+    api_keys: set[str] = set()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.load_api_keys()
+
+    def load_api_keys(self):
+        self.api_keys = set()
+        for env_var_name in _ENVIRONMENT_VARIABLE_NAMES:
+            if env_var := os.environ.get(env_var_name):
+                self.api_keys.update(set(key.strip() for key in env_var.split(",")))
+
+        if not self.api_keys:
+            raise MissingProjectKeyError()
+
+    def __get_random_api_key(self) -> str:
+        """
+        Get a random api key a private method.
+        """
+        return random.choice(list(self.api_keys))
 
     @property
     def uri(self) -> str:
@@ -42,15 +63,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}"
@@ -90,7 +103,14 @@ def connect(self):
         self._web3.eth.set_gas_price_strategy(rpc_gas_price_strategy)
 
     def disconnect(self):
+        """
+        Disconnect the connected API.
+        Refresh the API keys from environment variable.
+        Make the self.network_uris empty otherwise the old network_uri will be returned.
+        """
         self._web3 = None
+        self.load_api_keys()
+        self.network_uris = {}
 
     def get_virtual_machine_error(self, exception: Exception, **kwargs) -> VirtualMachineError:
         txn = kwargs.get("txn")
diff --git a/setup.py b/setup.py
index 0bc550e..ad8c31d 100644
--- a/setup.py
+++ b/setup.py
@@ -6,6 +6,7 @@
     "test": [  # `test` GitHub Action jobs uses this
         "pytest>=6.0",  # Core testing package
         "pytest-xdist",  # Multi-process runner
+        "pytest-mock",  # Mocking framework
         "pytest-cov",  # Coverage analyzer plugin
         "hypothesis>=6.2.0,<7.0",  # Strategy-based fuzzer
         "ape-arbitrum",  # For integration testing
diff --git a/tests/test_provider.py b/tests/test_provider.py
index ffa0c50..cd9f2aa 100644
--- a/tests/test_provider.py
+++ b/tests/test_provider.py
@@ -1,3 +1,5 @@
+import os
+
 import pytest
 import websocket  # type: ignore
 from ape.utils import ZERO_ADDRESS
@@ -31,3 +33,42 @@ 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, mocker):
+    mocker.patch.dict(
+        os.environ,
+        {"WEB3_INFURA_PROJECT_ID": "key1,key2,key3", "WEB3_INFURA_API_KEY": "key4,key5,key6"},
+    )
+    provider.load_api_keys()
+    # As there will be API keys in the ENV as well
+    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, mocker):
+    mocker.patch.dict(
+        os.environ,
+        {
+            "WEB3_INFURA_PROJECT_ID": "single_key1",
+            "WEB3_INFURA_API_KEY": "single_key2",
+        },
+    )
+    provider.load_api_keys()
+    assert len(provider.api_keys) == 2
+    assert "single_key1" in provider.api_keys
+    assert "single_key2" in provider.api_keys
+
+
+def test_uri_with_random_api_key(provider, mocker):
+    mocker.patch.dict(os.environ, {"WEB3_INFURA_PROJECT_ID": "key1, key2, key3, key4, key5, key6"})
+    provider.load_api_keys()
+    uris = set()
+    for _ in range(100):  # Generate multiple URIs
+        provider.disconnect()  # connect to a new URI
+        uri = provider.uri
+        uris.add(uri)
+        assert uri.startswith("https")
+        assert "/v3" in uri
+    assert len(uris) > 1  # Ensure we're getting different URIs with different