Feature: upgrade balance endpoint (#471)

Problem:
  Currently, the Balance API Endpoint only returns the total balance for an address. The user (also frontend) has no easy way to get the amount of locked tokens.

On : 
`/api/v0/addresses/{address}/balance`

Solution:
  Enhance the Balance API Endpoint to include the "locked_amount" field in the response, indicating the amount of tokens in use for a specific address.

Example:

```json
{
    "address": "0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba",
    "balance": 22192,
    "locked_amount": 2726
}
```

src/aleph/schemas/api/accounts.py

@@ -11,6 +11,7 @@
 class GetAccountBalanceResponse(BaseModel):
     address: str
     balance: Decimal
+    locked_amount: Decimal
 
 
 class GetAccountFilesQueryParams(BaseModel):

src/aleph/web/controllers/accounts.py

@@ -7,6 +7,7 @@
 from pydantic import ValidationError, parse_obj_as
 
 from aleph.db.accessors.balances import get_total_balance
+from aleph.db.accessors.cost import get_total_cost_for_address
 from aleph.db.accessors.files import (
     get_address_files_for_api,
     get_address_files_stats,
@@ -15,7 +16,8 @@
 from aleph.schemas.api.accounts import (
     GetAccountBalanceResponse,
     GetAccountFilesResponse,
-    GetAccountFilesQueryParams, GetAccountFilesResponseItem,
+    GetAccountFilesQueryParams,
+    GetAccountFilesResponseItem,
 )
 from aleph.types.db_session import DbSessionFactory
 from aleph.web.controllers.app_state_getters import get_session_factory_from_request
@@ -68,12 +70,15 @@ async def get_account_balance(request: web.Request):
         balance = get_total_balance(
             session=session, address=address, include_dapps=False
         )
+        total_cost = get_total_cost_for_address(session=session, address=address)
 
     if balance is None:
         raise web.HTTPNotFound()
 
     return web.json_response(
-        text=GetAccountBalanceResponse(address=address, balance=balance).json()
+        text=GetAccountBalanceResponse(
+            address=address, balance=balance, locked_amount=total_cost
+        ).json()
     )
 
 

tests/api/conftest.py

@@ -5,8 +5,10 @@
 import pytest
 import pytest_asyncio
 from aleph_message.models import AggregateContent, PostContent
+from configmanager import Config
 from sqlalchemy import insert
 
+from aleph.chains.chain_service import ChainService
 from aleph.db.accessors.aggregates import refresh_aggregate
 from aleph.db.models import (
     MessageDb,
@@ -15,10 +17,16 @@
     message_confirmations,
 )
 from aleph.db.models.posts import PostDb
+from aleph.handlers.message_handler import MessageHandler
+from aleph.jobs.process_pending_messages import PendingMessageProcessor
+from aleph.storage import StorageService
 from aleph.toolkit.timestamp import timestamp_to_datetime
 from aleph.types.db_session import DbSessionFactory
 import datetime as dt
 
+from in_memory_storage_engine import InMemoryStorageEngine
+
+
 # TODO: remove the raw parameter, it's just to avoid larger refactorings
 async def _load_fixtures(
     session_factory: DbSessionFactory, filename: str, raw: bool = True
@@ -33,7 +41,6 @@ async def _load_fixtures(
     tx_hashes = set()
 
     with session_factory() as session:
-
         for message_dict in messages_json:
             message_db = MessageDb.from_message_dict(message_dict)
             messages.append(message_db)
@@ -90,7 +97,7 @@ async def fixture_aggregate_messages(
             aggregate_keys.add((aggregate_element.owner, aggregate_element.key))
         session.commit()
 
-        for (owner, key) in aggregate_keys:
+        for owner, key in aggregate_keys:
             refresh_aggregate(session=session, owner=owner, key=key)
 
         session.commit()
@@ -188,3 +195,30 @@ def amended_post_with_refs_and_tags(post_with_refs_and_tags: Tuple[MessageDb, Po
     )
 
     return amend_message, amend_post
+
+
+@pytest.fixture
+def message_processor(mocker, mock_config: Config, session_factory: DbSessionFactory):
+    storage_engine = InMemoryStorageEngine(files={})
+    storage_service = StorageService(
+        storage_engine=storage_engine,
+        ipfs_service=mocker.AsyncMock(),
+        node_cache=mocker.AsyncMock(),
+    )
+    chain_service = ChainService(
+        session_factory=session_factory, storage_service=storage_service
+    )
+    message_handler = MessageHandler(
+        session_factory=session_factory,
+        chain_service=chain_service,
+        storage_service=storage_service,
+        config=mock_config,
+    )
+    message_processor = PendingMessageProcessor(
+        session_factory=session_factory,
+        message_handler=message_handler,
+        max_retries=0,
+        mq_message_exchange=mocker.AsyncMock(),
+        mq_conn=mocker.AsyncMock(),
+    )
+    return message_processor

tests/api/test_balance.py

@@ -0,0 +1,30 @@
+import json
+
+import pytest
+from aleph.db.models import (
+    AlephBalanceDb,
+)
+from aleph.jobs.process_pending_messages import PendingMessageProcessor
+
+MESSAGES_URI = "/api/v0/addresses/0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba/balance"
+
+
+@pytest.mark.asyncio
+async def test_get_balance(
+    ccn_api_client,
+    message_processor: PendingMessageProcessor,
+    instance_message_with_volumes_in_db,
+    fixture_instance_message,
+    user_balance: AlephBalanceDb,
+):
+    pipeline = message_processor.make_pipeline()
+    # Exhaust the iterator
+    _ = [message async for message in pipeline]
+
+    assert fixture_instance_message.item_content
+
+    response = await ccn_api_client.get(MESSAGES_URI)
+    assert response.status == 200, await response.text()
+    data = await response.json()
+    assert data["balance"] == user_balance.balance
+    assert data["locked_amount"] == 2726

tests/conftest.py

@@ -1,26 +1,50 @@
 import asyncio
 import contextlib
+import json
 import logging
 import os
 import shutil
 import sys
 from pathlib import Path
+from typing import Protocol, List
 
 import alembic.command
 import alembic.config
 import pytest
 import pytest_asyncio
+import pytz
+from aleph_message.models import (
+    MessageType,
+    Chain,
+    ItemType,
+    ExecutableContent,
+    ProgramContent,
+    InstanceContent,
+)
+from aleph_message.models.execution.volume import ImmutableVolume
 from configmanager import Config
 
 import aleph.config
+from aleph.db.accessors.files import insert_message_file_pin, upsert_file_tag
 from aleph.db.connection import make_engine, make_session_factory, make_db_url
+from aleph.db.models import (
+    PendingMessageDb,
+    MessageStatusDb,
+    StoredFileDb,
+    AlephBalanceDb,
+)
 from aleph.services.cache.node_cache import NodeCache
 from aleph.services.ipfs import IpfsService
 from aleph.services.ipfs.common import make_ipfs_client
 from aleph.services.storage.fileystem_engine import FileSystemStorageEngine
 from aleph.storage import StorageService
-from aleph.types.db_session import DbSessionFactory
+from aleph.toolkit.timestamp import timestamp_to_datetime
+from aleph.types.db_session import DbSessionFactory, DbSession
+from aleph.types.files import FileType, FileTag
+from aleph.types.message_status import MessageStatus
 from aleph.web import create_aiohttp_app
+from decimal import Decimal
+import datetime as dt
 
 # Add the helpers to the PYTHONPATH.
 # Note: mark the "helpers" directory as a source directory to tell PyCharm
@@ -126,3 +150,184 @@ async def ccn_api_client(
     client = await aiohttp_client(app)
 
     return client
+
+
+@pytest.fixture
+def fixture_instance_message(session_factory: DbSessionFactory) -> PendingMessageDb:
+    content = {
+        "address": "0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba",
+        "allow_amend": False,
+        "variables": {
+            "VM_CUSTOM_VARIABLE": "SOMETHING",
+            "VM_CUSTOM_VARIABLE_2": "32",
+        },
+        "environment": {
+            "reproducible": True,
+            "internet": False,
+            "aleph_api": False,
+            "shared_cache": False,
+        },
+        "resources": {"vcpus": 1, "memory": 128, "seconds": 30},
+        "requirements": {"cpu": {"architecture": "x86_64"}},
+        "rootfs": {
+            "parent": {
+                "ref": "549ec451d9b099cad112d4aaa2c00ac40fb6729a92ff252ff22eef0b5c3cb613",
+                "use_latest": True,
+            },
+            "persistence": "host",
+            "name": "test-rootfs",
+            "size_mib": 20 * 1024,
+        },
+        "authorized_keys": [
+            "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGULT6A41Msmw2KEu0R9MvUjhuWNAsbdeZ0DOwYbt4Qt user@example",
+            "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIH0jqdc5dmt75QhTrWqeHDV9xN8vxbgFyOYs2fuQl7CI",
+        ],
+        "volumes": [
+            {
+                "comment": "Python libraries. Read-only since a 'ref' is specified.",
+                "mount": "/opt/venv",
+                "ref": "5f31b0706f59404fad3d0bff97ef89ddf24da4761608ea0646329362c662ba51",
+                "use_latest": False,
+            },
+            {
+                "comment": "Ephemeral storage, read-write but will not persist after the VM stops",
+                "mount": "/var/cache",
+                "ephemeral": True,
+                "size_mib": 5,
+            },
+            {
+                "comment": "Working data persisted on the VM supervisor, not available on other nodes",
+                "mount": "/var/lib/sqlite",
+                "name": "sqlite-data",
+                "persistence": "host",
+                "size_mib": 10,
+            },
+            {
+                "comment": "Working data persisted on the Aleph network. "
+                "New VMs will try to use the latest version of this volume, "
+                "with no guarantee against conflicts",
+                "mount": "/var/lib/statistics",
+                "name": "statistics",
+                "persistence": "store",
+                "size_mib": 10,
+            },
+            {
+                "comment": "Raw drive to use by a process, do not mount it",
+                "name": "raw-data",
+                "persistence": "host",
+                "size_mib": 10,
+            },
+        ],
+        "time": 1619017773.8950517,
+    }
+
+    pending_message = PendingMessageDb(
+        item_hash="734a1287a2b7b5be060312ff5b05ad1bcf838950492e3428f2ac6437a1acad26",
+        type=MessageType.instance,
+        chain=Chain.ETH,
+        sender="0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba",
+        signature=None,
+        item_type=ItemType.inline,
+        item_content=json.dumps(content),
+        time=timestamp_to_datetime(1619017773.8950577),
+        channel=None,
+        reception_time=timestamp_to_datetime(1619017774),
+        fetched=True,
+        check_message=False,
+        retries=0,
+        next_attempt=dt.datetime(2023, 1, 1),
+    )
+    with session_factory() as session:
+        session.add(pending_message)
+        session.add(
+            MessageStatusDb(
+                item_hash=pending_message.item_hash,
+                status=MessageStatus.PENDING,
+                reception_time=pending_message.reception_time,
+            )
+        )
+        session.commit()
+
+    return pending_message
+
+
+@pytest.fixture
+def instance_message_with_volumes_in_db(
+    session_factory: DbSessionFactory, fixture_instance_message: PendingMessageDb
+) -> None:
+    with session_factory() as session:
+        insert_volume_refs(session, fixture_instance_message)
+        session.commit()
+
+
+class Volume(Protocol):
+    ref: str
+    use_latest: bool
+
+
+def get_volume_refs(content: ExecutableContent) -> List[Volume]:
+    volumes = []
+
+    for volume in content.volumes:
+        if isinstance(volume, ImmutableVolume):
+            volumes.append(volume)
+
+    if isinstance(content, ProgramContent):
+        volumes += [content.code, content.runtime]
+        if content.data:
+            volumes.append(content.data)
+
+    elif isinstance(content, InstanceContent):
+        if parent := content.rootfs.parent:
+            volumes.append(parent)
+
+    return volumes
+
+
+def insert_volume_refs(session: DbSession, message: PendingMessageDb):
+    """
+    Insert volume references in the DB to make the program processable.
+    """
+
+    content = InstanceContent.parse_raw(message.item_content)
+    volumes = get_volume_refs(content)
+
+    created = pytz.utc.localize(dt.datetime(2023, 1, 1))
+
+    for volume in volumes:
+        # Note: we use the reversed ref to generate the file hash for style points,
+        # but it could be set to any valid hash.
+        file_hash = volume.ref[::-1]
+
+        session.add(StoredFileDb(hash=file_hash, size=1024 * 1024, type=FileType.FILE))
+        session.flush()
+        insert_message_file_pin(
+            session=session,
+            file_hash=volume.ref[::-1],
+            owner=content.address,
+            item_hash=volume.ref,
+            ref=None,
+            created=created,
+        )
+        upsert_file_tag(
+            session=session,
+            tag=FileTag(volume.ref),
+            owner=content.address,
+            file_hash=volume.ref[::-1],
+            last_updated=created,
+        )
+
+
+@pytest.fixture
+def user_balance(session_factory: DbSessionFactory) -> AlephBalanceDb:
+    balance = AlephBalanceDb(
+        address="0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba",
+        chain=Chain.ETH,
+        balance=Decimal(22_192),
+        eth_height=0,
+    )
+
+    with session_factory() as session:
+        session.add(balance)
+        session.commit()
+    return balance