diff --git a/apps/dashboard_app/helpers/loans_table.py b/apps/dashboard_app/helpers/loans_table.py index c71a0178..21c6ffc1 100644 --- a/apps/dashboard_app/helpers/loans_table.py +++ b/apps/dashboard_app/helpers/loans_table.py @@ -1,3 +1,4 @@ + import pandas as pd from data_handler.handlers.loan_states.nostra_alpha.events import NostraAlphaState from data_handler.handlers.loan_states.nostra_mainnet.events import NostraMainnetState @@ -7,16 +8,7 @@ def get_protocol(state: State) -> str: - # TODO: Improve the inference. - if isinstance(state, ZkLendState): - return "zkLend" - if isinstance(state, NostraAlphaState) and not isinstance( - state, NostraMainnetState - ): - return "Nostra Alpha" - if isinstance(state, NostraMainnetState): - return "Nostra Mainnet" - raise ValueError + return state.get_protocol_name def get_loans_table_data( diff --git a/apps/dashboard_app/tests/test_state.py b/apps/dashboard_app/tests/test_state.py new file mode 100644 index 00000000..36a579f5 --- /dev/null +++ b/apps/dashboard_app/tests/test_state.py @@ -0,0 +1,43 @@ +import pytest +from shared.state import State + + +class MockLoanEntity: + """Mock loan entity for testing""" + pass + + +class MockState(State): + def __init__(self): + super().__init__(loan_entity_class=MockLoanEntity) + + def compute_liquidable_debt_at_price(self, *args, **kwargs): + return None + + +def test_state_without_protocol_name(): + """Test that accessing get_protocol_name raises NotImplementedError when PROTOCOL_NAME is not set""" + mock_state = MockState() + with pytest.raises(NotImplementedError): + _ = mock_state.get_protocol_name + + +@pytest.mark.parametrize( + "protocol_name", + ["zkLend", "Nostra Alpha", "Nostra Mainnet"] +) +def test_protocol_names(protocol_name): + """Test that protocol name is correctly returned""" + mock_state = MockState() + mock_state.PROTOCOL_NAME = protocol_name + assert mock_state.get_protocol_name == protocol_name + + +def test_get_protocol_helper(): + """Test the get_protocol helper function""" + from dashboard_app.helpers.loans_table import get_protocol + + mock_state = MockState() + mock_state.PROTOCOL_NAME = "Test Protocol" + assert get_protocol(mock_state) == "Test Protocol" + \ No newline at end of file diff --git a/apps/data_handler/handlers/loan_states/hashtack_v0/events.py b/apps/data_handler/handlers/loan_states/hashtack_v0/events.py index c1c493d9..4f076e8a 100644 --- a/apps/data_handler/handlers/loan_states/hashtack_v0/events.py +++ b/apps/data_handler/handlers/loan_states/hashtack_v0/events.py @@ -14,7 +14,7 @@ from data_handler.handlers.helpers import get_symbol from data_handler.db.crud import InitializerDBConnector -from shared.constants import TOKEN_SETTINGS +from shared.constants import ProtocolIDs, TOKEN_SETTINGS from shared.loan_entity import LoanEntity from shared.state import State from shared.types import InterestRateModels, Portfolio, TokenValues, TokenSettings @@ -184,6 +184,7 @@ class HashstackV0State(State): debt, thus we always rewrite the balances whenever they are updated. """ + PROTOCOL_NAME: str = ProtocolIDs.HASHSTACK_V0.value EVENTS_METHODS_MAPPING: dict[str, str] = EVENTS_METHODS_MAPPING def __init__( diff --git a/apps/data_handler/handlers/loan_states/hashtack_v1/events.py b/apps/data_handler/handlers/loan_states/hashtack_v1/events.py index 6edb5ec6..d74b4472 100644 --- a/apps/data_handler/handlers/loan_states/hashtack_v1/events.py +++ b/apps/data_handler/handlers/loan_states/hashtack_v1/events.py @@ -12,7 +12,7 @@ from data_handler.handlers.settings import TokenSettings from data_handler.db.crud import InitializerDBConnector -from shared.constants import TOKEN_SETTINGS +from shared.constants import TOKEN_SETTINGS, ProtocolIDs from shared.helpers import add_leading_zeros from shared.loan_entity import LoanEntity from shared.state import State @@ -431,6 +431,7 @@ class HashstackV1State(State): debt, thus we always rewrite the balances whenever they are updated. """ + PROTOCOL_NAME: str = ProtocolIDs.HASHSTACK_V1.value ADDRESSES_TO_TOKENS: dict[str, str] = ADDRESSES_TO_TOKENS EVENTS_METHODS_MAPPING: dict[str, str] = EVENTS_METHODS_MAPPING diff --git a/apps/data_handler/handlers/loan_states/nostra_mainnet/events.py b/apps/data_handler/handlers/loan_states/nostra_mainnet/events.py index 55175604..24f9c8b7 100644 --- a/apps/data_handler/handlers/loan_states/nostra_mainnet/events.py +++ b/apps/data_handler/handlers/loan_states/nostra_mainnet/events.py @@ -113,6 +113,7 @@ class NostraMainnetState(NostraAlphaState): relevant event are implemented in `.nostra_alpha.NostraAlphaState`. """ + PROTOCOL_NAME: str = ProtocolIDs.NOSTRA_MAINNET.value TOKEN_ADDRESSES: list[str] = NOSTRA_MAINNET_TOKEN_ADDRESSES INTEREST_RATE_MODEL_ADDRESS: str = NOSTRA_MAINNET_INTEREST_RATE_MODEL_ADDRESS CDP_MANAGER_ADDRESS: str = NOSTRA_MAINNET_CDP_MANAGER_ADDRESS diff --git a/apps/data_handler/handlers/loan_states/zklend/events.py b/apps/data_handler/handlers/loan_states/zklend/events.py index 008fdb17..8c9e9e30 100644 --- a/apps/data_handler/handlers/loan_states/zklend/events.py +++ b/apps/data_handler/handlers/loan_states/zklend/events.py @@ -44,6 +44,7 @@ ZkLendCollateralTokenParameters, ZkLendDebtTokenParameters, ) +from shared.constants import ProtocolIDs ZKLEND_MARKET: str = ("0x04c0a5193d58f74fbace4b74dcf65481e734ed1714121bdc571da345540efa05") EVENTS_METHODS_MAPPING: dict[str, str] = { @@ -166,6 +167,7 @@ class ZkLendState(State): relevant event. """ + PROTOCOL_NAME: str = ProtocolIDs.ZKLEND.value EVENTS_METHODS_MAPPING: dict[str, str] = EVENTS_METHODS_MAPPING def __init__( diff --git a/apps/data_handler/tests/health_ratio_test/test_health_ratio_tests.py b/apps/data_handler/tests/health_ratio_test/test_health_ratio_tests.py index 5b3a9597..05f60c7c 100644 --- a/apps/data_handler/tests/health_ratio_test/test_health_ratio_tests.py +++ b/apps/data_handler/tests/health_ratio_test/test_health_ratio_tests.py @@ -16,6 +16,8 @@ class MockState(State): """Mock State class for testing""" + PROTOCOL_NAME: str = ProtocolIDs.ZKLEND.value + def __init__(self): self.loan_entities = {} self.collateral_interest_rate_models = None diff --git a/apps/shared/state.py b/apps/shared/state.py index 5a1bb6e2..101df551 100644 --- a/apps/shared/state.py +++ b/apps/shared/state.py @@ -43,6 +43,13 @@ def __init__( self.last_block_number: int = 0 self.last_interest_rate_block_number: int = 0 + @property + def get_protocol_name(self) -> str: + """Returns the protocol name for the state""" + if not self.PROTOCOL_NAME: + raise NotImplementedError("PROTOCOL_NAME must be set in the implementing class") + return self.PROTOCOL_NAME + def process_event(self, method_name: str, event: pd.Series) -> None: # TODO: Save the timestamp of each update? if event["block_number"] >= self.last_block_number: diff --git a/apps/web_app/utils/zklend.py b/apps/web_app/utils/zklend.py index 9c7c5b7c..7f992fc8 100644 --- a/apps/web_app/utils/zklend.py +++ b/apps/web_app/utils/zklend.py @@ -5,6 +5,7 @@ from .settings import TOKEN_SETTINGS as BASE_TOKEN_SETTINGS from .settings import TokenSettings as BaseTokenSettings from .state import InterestRateModels, LoanEntity, State +from shared.constants import ProtocolIDs # IT GIVES `ModuleNotFoundError` THAT'S WHY I COMMENTED OUT IT # from data_handler.handlers.loan_states.zklend.fetch_zklend_specific_token_settings import ZKLEND_SPECIFIC_TOKEN_SETTINGS @@ -124,6 +125,7 @@ class ZkLendState(State): relevant event. """ + PROTOCOL_NAME: str = ProtocolIDs.ZKLEND.value EVENTS_METHODS_MAPPING: dict[str, str] = EVENTS_METHODS_MAPPING def __init__(