From 7bc1d828983de57431c9bce641395bcddb2d9dce Mon Sep 17 00:00:00 2001
From: Kahtaf Alam <kahtaf@gmail.com>
Date: Tue, 3 Dec 2024 17:15:37 -0500
Subject: [PATCH] Allow overriding default contract addresses (#50)

---
 pyproject.toml             |  2 +-
 vana/__init__.py           |  2 +-
 vana/client.py             | 32 +++++++++++++++++++++++++++-----
 vana/contracts/__init__.py |  4 ++--
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 893a9b1..e36710b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vana"
-version = "0.32.0"
+version = "0.33.0"
 description = ""
 authors = ["Tim Nunamaker <tim@vana.com>", "Volodymyr Isai <volod@vana.com>", "Kahtaf Alam <kahtaf@vana.com>"]
 readme = "README.md"
diff --git a/vana/__init__.py b/vana/__init__.py
index 17c340c..a5fcdce 100644
--- a/vana/__init__.py
+++ b/vana/__init__.py
@@ -15,7 +15,7 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-__version__ = "0.32.0"
+__version__ = "0.33.0"
 
 import rich
 
diff --git a/vana/client.py b/vana/client.py
index fbb7bb4..ff719e7 100644
--- a/vana/client.py
+++ b/vana/client.py
@@ -2,9 +2,8 @@
 import copy
 import json
 import os
-from typing import Optional
-
 import vana
+from typing import Optional
 from vana.chain_data import Proof, ProofData
 from vana.contracts import contracts
 from vana.utils.web3 import as_wad
@@ -14,9 +13,32 @@ class Client:
     @staticmethod
     def config() -> "config":
         parser = argparse.ArgumentParser()
-        vana.ChainManager.add_args(parser)
+        vana.Client.add_args(parser)
         return vana.Config(parser, args=[])
 
+    @classmethod
+    def add_args(cls, parser: argparse.ArgumentParser, prefix: Optional[str] = None):
+        prefix_str = "" if prefix is None else f"{prefix}."
+        try:
+            parser.add_argument(
+                "--" + prefix_str + "client.tee_pool_contract_address",
+                default=os.getenv("TEE_POOL_CONTRACT_ADDRESS") or None,
+                type=str,
+                help="""The address for the TEE Pool Contract.""")
+            parser.add_argument(
+                "--" + prefix_str + "client.data_registry_contract_address",
+                default=os.getenv("DATA_REGISTRY_CONTRACT_ADDRESS") or None,
+                type=str,
+                help="""The address for the Data Registry Contract.""")
+            parser.add_argument(
+                "--" + prefix_str + "client.dlp_root_contract_address",
+                default=os.getenv("DLP_ROOT_CONTRACT_ADDRESS") or None,
+                type=str,
+                help="""The address for the DLP Root Contract.""")
+        except argparse.ArgumentError:
+            # re-parsing arguments.
+            pass
+
     def __init__(self, config: vana.Config):
         if config is None:
             config = self.config()
@@ -33,7 +55,7 @@ def __init__(self, config: vana.Config):
         )
         with open(data_registry_contract_path) as f:
             self.data_registry_contract = self.chain_manager.web3.eth.contract(
-                address=contracts[self.network]["DataRegistry"],
+                address=config.client.data_registry_contract_address or contracts[self.network]["DataRegistry"],
                 abi=json.load(f)
             )
         tee_pool_contract_path = os.path.join(
@@ -42,7 +64,7 @@ def __init__(self, config: vana.Config):
         )
         with open(tee_pool_contract_path) as f:
             self.tee_pool_contract = self.chain_manager.web3.eth.contract(
-                address=contracts[self.network]["TeePool"],
+                address=config.client.tee_pool_contract_address or contracts[self.network]["TeePool"],
                 abi=json.load(f)
             )
 
diff --git a/vana/contracts/__init__.py b/vana/contracts/__init__.py
index 7bb6bfc..65e878d 100644
--- a/vana/contracts/__init__.py
+++ b/vana/contracts/__init__.py
@@ -2,12 +2,12 @@
     "vana": {
         "TeePool": "0x3c92fD91639b41f13338CE62f19131e7d19eaa0D",
         "DataRegistry": "0x8C8788f98385F6ba1adD4234e551ABba0f82Cb7C",
-        "RootNetworkContract": "0xf408A064d640b620219F510963646Ed2bD5606BB",
+        "RootNetworkContract": "0xff14346dF2B8Fd0c95BF34f1c92e49417b508AD5",
     },
     "islander": {
         "TeePool": "0x3c92fD91639b41f13338CE62f19131e7d19eaa0D",
         "DataRegistry": "0x8C8788f98385F6ba1adD4234e551ABba0f82Cb7C",
-        "RootNetworkContract": "0xf408A064d640b620219F510963646Ed2bD5606BB",
+        "RootNetworkContract": "0xff14346dF2B8Fd0c95BF34f1c92e49417b508AD5",
     },
     "maya": {
         "TeePool": "0xF084Ca24B4E29Aa843898e0B12c465fAFD089965",