From 74ca0752f0ae47be5edc7d19d5faccfa54d120e5 Mon Sep 17 00:00:00 2001
From: Juliya Smith <jules@apeworx.io>
Date: Mon, 25 Nov 2024 17:23:50 -0600
Subject: [PATCH] feat: add remaining networks

---
 README.md              | 11 ++++++++++-
 ape_infura/provider.py | 22 +++++++++++++++-------
 ape_infura/utils.py    | 40 +++++++++++++++++++++++++++++++++++-----
 tests/test_provider.py |  2 +-
 4 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index d3032b9..9a9e953 100644
--- a/README.md
+++ b/README.md
@@ -5,11 +5,20 @@ This plugin supports the following ecosystems:
 
 - Arbitrum
 - Avalanche
+- Base
 - Blast
+- BSC
+- Celo
 - Ethereum
-- ~~Linea~~ (awaiting ape-linea update)
+- Linea
+- Mantle
 - Optimism
+- Palm
 - Polygon
+- Scroll
+- Starknet
+- Unichain
+- zkSync
 
 ## Dependencies
 
diff --git a/ape_infura/provider.py b/ape_infura/provider.py
index ca5515c..6cfd6b6 100644
--- a/ape_infura/provider.py
+++ b/ape_infura/provider.py
@@ -13,12 +13,18 @@
 
 _ENVIRONMENT_VARIABLE_NAMES = ("WEB3_INFURA_PROJECT_ID", "WEB3_INFURA_API_KEY")
 # NOTE: https://docs.infura.io/learn/websockets#supported-networks
-_WEBSOCKET_CAPABLE_ECOSYSTEMS = {
-    "ethereum",
-    "arbitrum",
-    "optimism",
-    "polygon",
-    "linea",
+_WEBSOCKET_CAPABLE_NETWORKS = {
+    "arbitrum": ("mainnet", "sepolia"),
+    "avalanche": ("fuji", "mainnet"),
+    "base": ("mainnet", "sepolia"),
+    "blast": ("mainnet",),
+    "bsc": ("mainnet", "opbnb"),
+    "ethereum": ("holesky", "mainnet", "sepolia"),
+    "linea": ("mainnet", "sepolia"),
+    "mainnet": ("mainnet",),
+    "optimism": ("mainnet", "sepolia"),
+    "polygon": ("amoy", "mainnet"),
+    "scroll": ("mainnet",),
 }
 
 
@@ -83,7 +89,9 @@ def http_uri(self) -> str:
     @property
     def ws_uri(self) -> Optional[str]:
         # NOTE: Overriding `Web3Provider.ws_uri` implementation
-        if self.network.ecosystem.name not in _WEBSOCKET_CAPABLE_ECOSYSTEMS:
+        ecosystem_name = self.network.ecosystem.name
+        network_name = self.network.name
+        if network_name not in _WEBSOCKET_CAPABLE_NETWORKS.get(ecosystem_name, []):
             return None
 
         # Remove `http` in default URI w/ `ws`, also infura adds `/ws` to URI
diff --git a/ape_infura/utils.py b/ape_infura/utils.py
index 35a4fda..7a815a1 100644
--- a/ape_infura/utils.py
+++ b/ape_infura/utils.py
@@ -7,24 +7,54 @@
         "fuji",
         "mainnet",
     ],
+    "base": [
+        "mainnet",
+        "sepolia",
+    ],
     "blast": [
         "mainnet",
         "sepolia",
     ],
+    "bsc": ["mainnet", "opbnb", "opbnb-testnet", "testnet"],
+    "celo": [
+        "alfajores",
+        "mainnet",
+    ],
     "ethereum": [
+        "holesky",
+        "mainnet",
+        "sepolia",
+    ],
+    "linea": [
+        "mainnet",
+        "sepolia",
+    ],
+    "mantle": [
         "mainnet",
         "sepolia",
     ],
-    # TODO: Comment out after ape-linea supports 0.7
-    # "linea": [
-    #     "mainnet",
-    # ],
     "optimism": [
         "mainnet",
         "sepolia",
     ],
-    "polygon": [
+    "palm": [
         "mainnet",
+        "testnet",
+    ],
+    "polygon": [
         "amoy",
+        "mainnet",
+    ],
+    "scroll": ["mainnet", "sepolia"],
+    "starknet": [
+        "mainnet",
+        "sepolia",
+    ],
+    "unichain": [
+        "sepolia",
+    ],
+    "zksync": [
+        "mainnet",
+        "sepolia",
     ],
 }
diff --git a/tests/test_provider.py b/tests/test_provider.py
index bddbe28..d61c489 100644
--- a/tests/test_provider.py
+++ b/tests/test_provider.py
@@ -4,7 +4,7 @@
 import websocket  # type: ignore
 from ape.utils import ZERO_ADDRESS
 
-from ape_infura.provider import _WEBSOCKET_CAPABLE_ECOSYSTEMS, Infura
+from ape_infura.provider import _WEBSOCKET_CAPABLE_NETWORKS, Infura
 
 
 def test_infura_http(provider):