diff --git a/src/aleph/vm/orchestrator/payment.py b/src/aleph/vm/orchestrator/payment.py index cfa2e529..7de0c51b 100644 --- a/src/aleph/vm/orchestrator/payment.py +++ b/src/aleph/vm/orchestrator/payment.py @@ -40,8 +40,12 @@ async def get_address_balance(address: str) -> dict: # Raise an error if the request failed resp.raise_for_status() + ret = await resp.json() + if not isinstance(ret, dict): + msg = f"get_address_blaance: Invalid response received from {url}, should be dict" + raise Exception(msg) - return await resp.json() + return ret async def fetch_balance_of_address(address: str) -> Decimal: @@ -50,10 +54,13 @@ async def fetch_balance_of_address(address: str) -> Decimal: """ resp_data = await get_address_balance(address) - if hasattr(resp_data, "balance"): - return Decimal(resp_data["balance"]) - - return Decimal(0) + balance = resp_data.get("balance", 0) + if not isinstance(balance, Decimal | int | float | str | tuple): + logger.warning( + "Invalid balance format for address %s: %r. API Server : %s", address, balance, settings.API_SERVER + ) + return Decimal(0) + return Decimal(balance) async def fetch_credit_balance_of_address(address: str) -> Decimal: @@ -62,10 +69,13 @@ async def fetch_credit_balance_of_address(address: str) -> Decimal: """ resp_data = await get_address_balance(address) - if hasattr(resp_data, "credit_balance"): - return Decimal(resp_data["credit_balance"]) - - return Decimal(0) + balance = resp_data.get("credit_balance", 0) + if not isinstance(balance, Decimal | int | float | str | tuple): + logger.warning( + "Invalid balance format for address %s: %r. API Server : %s", address, balance, settings.API_SERVER + ) + return Decimal(0) + return Decimal(balance) async def fetch_execution_price( diff --git a/tests/supervisor/test_payment.py b/tests/supervisor/test_payment.py new file mode 100644 index 00000000..c5cd3369 --- /dev/null +++ b/tests/supervisor/test_payment.py @@ -0,0 +1,64 @@ +from decimal import Decimal + +import pytest + +from aleph.vm.conf import settings +from aleph.vm.orchestrator.payment import ( + fetch_balance_of_address, + fetch_credit_balance_of_address, +) + + +@pytest.fixture() +def mock_get_address_balance(mocker): + # http://51.159.106.166:4024/api/v0/addresses/0x5f78199cd833c1dc1735bee4a7416caaE58Facca/balance + fake = { + "address": "0x555559cd833c1dc1735bee4a7416caaE58Facca", + "balance": 57000.0, + "details": {"AVAX": 3000.0, "BASE": 52000.0, "ETH": 2000.0}, + "locked_amount": 4010.008710650127, + "credit_balance": 10000, + } + mocker.patch("aleph.vm.orchestrator.payment.get_address_balance", return_value=fake) + + return fake + + +@pytest.mark.asyncio +async def test_fetch_balance_of_address(mocker, mock_get_address_balance): + """ """ + mocker.patch.object(settings, "API_SERVER", "https://fake.aleph.cloud") + balance = await fetch_balance_of_address("0x555559cd833c1dc1735bee4a7416caaE58Facca") + assert balance == Decimal("57000") + + +@pytest.mark.asyncio +async def test_fetch_credit_balance_of_address(mocker, mock_get_address_balance): + """ """ + mocker.patch.object(settings, "API_SERVER", "https://fake.aleph.cloud") + balance = await fetch_credit_balance_of_address("0x555559cd833c1dc1735bee4a7416caaE58Facca") + assert balance == Decimal("10000.0") + + +@pytest.mark.asyncio +async def test_fetch_credit_balance_of_address_empty_response( + mocker, +): + """ """ + mocker.patch.object(settings, "API_SERVER", "https://fake.aleph.cloud") + mocker.patch("aleph.vm.orchestrator.payment.get_address_balance", return_value={}) + + balance = await fetch_credit_balance_of_address("0x555559cd833c1dc1735bee4a7416caaE58Facca") + assert balance == Decimal("0") + + +@pytest.mark.asyncio +async def test_fetch_balance_of_address_empty_response( + mocker, +): + """ """ + mocker.patch.object(settings, "API_SERVER", "https://fake.aleph.cloud") + mocker.patch("aleph.vm.orchestrator.payment.get_address_balance", return_value={}) + + balance = await fetch_balance_of_address("0x555559cd833c1dc1735bee4a7416caaE58Facca") + assert balance == Decimal("0")