From f54e723dafeed08c355700dfd031088337719365 Mon Sep 17 00:00:00 2001
From: niv vaknin <niv@certora.com>
Date: Sun, 5 Jan 2025 11:08:01 +0200
Subject: [PATCH] Reverse mapping

---
 Quorum/apis/governance/aave_governance.py | 18 +++++++++++++++++-
 Quorum/apis/governance/data_models.py     |  4 +++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/Quorum/apis/governance/aave_governance.py b/Quorum/apis/governance/aave_governance.py
index d272520..29b9d40 100644
--- a/Quorum/apis/governance/aave_governance.py
+++ b/Quorum/apis/governance/aave_governance.py
@@ -1,9 +1,25 @@
 import requests
+
+from Quorum.utils.chain_enum import Chain
 from Quorum.apis.governance.data_models import BGDProposalData, PayloadAddresses
 
 BASE_BGD_CACHE_REPO = 'https://raw.githubusercontent.com/bgd-labs/v3-governance-cache/refs/heads/main/cache'
 PROPOSALS_URL = f'{BASE_BGD_CACHE_REPO}/1/0x9AEE0B04504CeF83A65AC3f0e838D0593BCb2BC7/proposals'
 
+CHAIN_ID_TO_CHAIN = {
+    1: Chain.ETH,
+    42161: Chain.ARB,
+    43114: Chain.AVAX,
+    8453: Chain.BASE,
+    56: Chain.BSC,
+    100: Chain.GNO,
+    10: Chain.OPT,
+    137: Chain.POLY,
+    534352: Chain.SCROLL,
+    324: Chain.ZK,
+    59144: Chain.LINEA,
+}
+
 class AaveGovernanceAPI:
     """
     A utility class to interact with the BGD governance cache and retrieve 
@@ -65,5 +81,5 @@ def get_all_payloads_addresses(self, proposal_id: int) -> list[PayloadAddresses]
         results = []
         for p in data.proposal.payloads:
             addresses = self.get_payload_addresses(p.chain, p.payloads_controller, p.payload_id)
-            results.append(PayloadAddresses(chain=p.chain, addresses=addresses))
+            results.append(PayloadAddresses(chain=CHAIN_ID_TO_CHAIN[p.chain], addresses=addresses))
         return results
diff --git a/Quorum/apis/governance/data_models.py b/Quorum/apis/governance/data_models.py
index 6066933..72536d2 100644
--- a/Quorum/apis/governance/data_models.py
+++ b/Quorum/apis/governance/data_models.py
@@ -1,6 +1,8 @@
 from typing import List, Optional
 from pydantic import BaseModel, Field
 
+from Quorum.utils.chain_enum import Chain
+
 class IPFSData(BaseModel):
     title: Optional[str] = None
     discussions: Optional[str] = None
@@ -43,5 +45,5 @@ class BGDProposalData(BaseModel):
 
 
 class PayloadAddresses(BaseModel):
-    chain: str
+    chain: Chain
     addresses: List[str]