From 02133bfeae10e3b681cd4abce36cf951b9093aae Mon Sep 17 00:00:00 2001
From: Anton Rudenko <mrroudyk09@gmail.com>
Date: Thu, 5 Dec 2024 19:56:35 +0200
Subject: [PATCH] Fix token issues

---
 apps/dashboard_app/helpers/load_data.py      |  8 ++++---
 apps/dashboard_app/helpers/protocol_stats.py | 14 +++++++++---
 apps/dashboard_app/helpers/tools.py          | 23 +++++++-------------
 3 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/apps/dashboard_app/helpers/load_data.py b/apps/dashboard_app/helpers/load_data.py
index 9a1b28a9..682315c3 100644
--- a/apps/dashboard_app/helpers/load_data.py
+++ b/apps/dashboard_app/helpers/load_data.py
@@ -20,7 +20,7 @@
     get_supply_stats,
     get_utilization_stats,
 )
-from dashboard_app.helpers.tools import get_prices
+from dashboard_app.helpers.tools import add_leading_zeros, get_prices
 
 logger = logging.getLogger(__name__)
 
@@ -54,7 +54,9 @@ def _init_zklend_state(self) -> ZkLendState:
         logger.info("Initializing ZkLend state.")
         zklend_state = ZkLendState()
         start = monotonic()
-        zklend_data = self.data_connector.fetch_data(self.data_connector.ZKLEND_SQL_QUERY)
+        zklend_data = self.data_connector.fetch_data(
+            self.data_connector.ZKLEND_SQL_QUERY
+        )
         zklend_interest_rate_data = self.data_connector.fetch_data(
             self.data_connector.ZKLEND_INTEREST_RATE_SQL_QUERY
         )
@@ -114,7 +116,7 @@ def _set_underlying_addresses_to_decimals(self):
             )
         self.underlying_addresses_to_decimals.update(
             {
-                x.address: int(math.log10(x.decimal_factor))
+                add_leading_zeros(x.address): int(math.log10(x.decimal_factor))
                 for x in TOKEN_SETTINGS.values()
             }
         )
diff --git a/apps/dashboard_app/helpers/protocol_stats.py b/apps/dashboard_app/helpers/protocol_stats.py
index b79cfadf..e7921fe6 100644
--- a/apps/dashboard_app/helpers/protocol_stats.py
+++ b/apps/dashboard_app/helpers/protocol_stats.py
@@ -16,7 +16,11 @@
     get_protocol,
     get_supply_function_call_parameters,
 )
-from dashboard_app.helpers.tools import get_addresses, get_underlying_address
+from dashboard_app.helpers.tools import (
+    add_leading_zeros,
+    get_addresses,
+    get_underlying_address,
+)
 
 
 def get_general_stats(
@@ -116,7 +120,11 @@ def get_supply_stats(
     df = pd.DataFrame(data)
     df["Total supply (USD)"] = sum(
         df[column]
-        * Decimal(prices[TOKEN_SETTINGS[column.replace(" supply", "")].address])
+        * Decimal(
+            prices[
+                add_leading_zeros(TOKEN_SETTINGS[column.replace(" supply", "")].address)
+            ]
+        )
         for column in df.columns
         if "supply" in column
     ).apply(lambda x: round(x, 4))
@@ -248,7 +256,7 @@ def get_utilization_stats(
 ) -> pd.DataFrame:
     """
     Get utilization stats for the dashboard.
-    :param stats: DataFrame containing 
+    :param stats: DataFrame containing
     general_stats, supply_stats, debt_stat.
     :return: DataFrame with utilization stats
     """
diff --git a/apps/dashboard_app/helpers/tools.py b/apps/dashboard_app/helpers/tools.py
index 376bd707..52016756 100644
--- a/apps/dashboard_app/helpers/tools.py
+++ b/apps/dashboard_app/helpers/tools.py
@@ -40,7 +40,7 @@ def get_collateral_token_range(
     collateral_token_price: float,
 ) -> list[float]:
     """
-    Generates a range of prices for a collateral token and 
+    Generates a range of prices for a collateral token and
     Returns:  A list of float values representing the range of prices for the collateral token.
     """
     target_number_of_values = 50
@@ -54,7 +54,7 @@ def get_collateral_token_range(
     difference = [
         abs(50 - stop_price / (k * magnitude)) for k in step_factors
     ]  # Stores the difference between the target value and
-       #number of values generated from each step factor.
+    # number of values generated from each step factor.
     readable_step = (
         step_factors[difference.index(min(difference))] * magnitude
     )  # Gets readable step from step factor with values closest to the target value.
@@ -98,19 +98,9 @@ def get_prices(token_decimals: dict[str, int]) -> dict[str, float]:
 
     tokens_info = response.json()
 
-    # Define the addresses for which you do not want to apply add_leading_zeros
-    skip_leading_zeros_addresses = {
-        TOKEN_SETTINGS["STRK"].address,
-    }
-
     # Create a map of token addresses to token information, applying add_leading_zeros conditionally
     token_info_map = {
-        (
-            token["address"]
-            if token["address"] in skip_leading_zeros_addresses
-            else add_leading_zeros(token["address"])
-        ): token
-        for token in tokens_info
+        add_leading_zeros(token["address"]): token for token in tokens_info
     }
 
     prices = {}
@@ -123,8 +113,11 @@ def get_prices(token_decimals: dict[str, int]) -> dict[str, float]:
 
         if decimals != token_info.get("decimals"):
             logging.error(
-                "Decimal mismatch for token %s: expected %d, got %d"
-                 ,token, decimals, token_info.get('decimals'))
+                "Decimal mismatch for token %s: expected %d, got %d",
+                token,
+                decimals,
+                token_info.get("decimals"),
+            )
             continue
 
         prices[token] = token_info.get("currentPrice")