diff --git a/quantara/web_app/alembic/versions/b2c3d4e5f6a7_add_vault_unique_constraint.py b/quantara/web_app/alembic/versions/b2c3d4e5f6a7_add_vault_unique_constraint.py new file mode 100644 index 00000000..b5438526 --- /dev/null +++ b/quantara/web_app/alembic/versions/b2c3d4e5f6a7_add_vault_unique_constraint.py @@ -0,0 +1,73 @@ +"""Add unique constraint on vault(user_id, symbol) + +Revision ID: b2c3d4e5f6a7 +Revises: outbox_event_table_rev +Create Date: 2026-08-20 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "b2c3d4e5f6a7" +down_revision = "outbox_event_table_rev" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + """Deduplicates vault rows then adds a unique constraint on (user_id, symbol).""" + conn = op.get_bind() + + # Deduplicate: for each (user_id, symbol) pair keep only the row with the + # latest updated_at and sum all amounts into it, then delete the rest. + conn.execute(sa.text(""" + WITH ranked AS ( + SELECT id, + user_id, + symbol, + amount, + updated_at, + ROW_NUMBER() OVER ( + PARTITION BY user_id, symbol ORDER BY updated_at DESC, id + ) AS rn + FROM vault + ), + totals AS ( + SELECT user_id, + symbol, + SUM(amount::numeric) AS total_amount + FROM vault + GROUP BY user_id, symbol + ) + UPDATE vault v + SET amount = t.total_amount::text + FROM totals t + WHERE v.user_id = t.user_id + AND v.symbol = t.symbol + AND v.id IN ( + SELECT id FROM ranked WHERE rn > 1 + ); + + DELETE FROM vault + WHERE id IN ( + SELECT id FROM ( + SELECT id, + ROW_NUMBER() OVER ( + PARTITION BY user_id, symbol ORDER BY updated_at DESC, id + ) AS rn + FROM vault + ) sub + WHERE rn > 1 + ); + """)) + + op.create_unique_constraint( + "uq_vault_user_symbol", "vault", ["user_id", "symbol"] + ) + + +def downgrade() -> None: + """Drops the unique constraint on vault(user_id, symbol).""" + op.drop_constraint("uq_vault_user_symbol", "vault", type_="unique") diff --git a/quantara/web_app/db/crud/deposit.py b/quantara/web_app/db/crud/deposit.py index b6450b73..d021a8d0 100644 --- a/quantara/web_app/db/crud/deposit.py +++ b/quantara/web_app/db/crud/deposit.py @@ -3,9 +3,13 @@ """ import logging +import uuid from decimal import Decimal from typing import TypeVar +from sqlalchemy import Numeric, cast, func +from sqlalchemy.dialects.postgresql import insert as pg_insert + from web_app.db.models import Base, User, Vault from .base import DBConnector @@ -20,9 +24,42 @@ class DepositDBConnector(DBConnector): Provides database connection and operations management for the Vault model. """ + def upsert_vault(self, user_id: uuid.UUID, symbol: str, amount: str) -> Vault: + """ + Atomically inserts a new vault row or adds to the existing balance. + + Uses PostgreSQL INSERT ... ON CONFLICT DO UPDATE so that concurrent + calls for the same (user_id, symbol) never lose updates. + + :param user_id: UUID of the user + :param symbol: Token symbol or address + :param amount: Amount to add (as string) + + :return: Vault instance (existing row updated or newly inserted) + """ + with self.Session() as db: + stmt = ( + pg_insert(Vault) + .values(user_id=user_id, symbol=symbol, amount=amount) + .on_conflict_do_update( + constraint="uq_vault_user_symbol", + set_={ + "amount": cast(Vault.amount, Numeric) + + cast(amount, Numeric), + "updated_at": func.now(), + }, + ) + .returning(Vault) + ) + result = db.execute(stmt) + vault = result.scalar_one() + db.commit() + db.refresh(vault) + return vault + def create_vault(self, user: User, symbol: str, amount: str) -> Vault: """ - Creates a new vault instance + Creates a new vault instance or updates existing balance atomically. :param user: A user model instance :param symbol: Token symbol or address @@ -30,9 +67,7 @@ def create_vault(self, user: User, symbol: str, amount: str) -> Vault: :return: Vault """ - vault = Vault(user_id=user.id, symbol=symbol, amount=amount) - self.write_to_db(vault) - return vault + return self.upsert_vault(user.id, symbol, amount) def get_vault(self, wallet_id: str, symbol: str) -> Vault | None: """ @@ -53,7 +88,7 @@ def get_vault(self, wallet_id: str, symbol: str) -> Vault | None: def add_vault_balance(self, wallet_id: str, symbol: str, amount: str) -> Vault: """ - Adds balance to user vault for token symbol + Adds balance to user vault for token symbol atomically. :param wallet_id: Wallet id of user :param symbol: Token symbol or address @@ -61,15 +96,10 @@ def add_vault_balance(self, wallet_id: str, symbol: str, amount: str) -> Vault: :return: Updated Vault instance """ - vault = self.get_vault(wallet_id, symbol) - if not vault: - raise ValueError("Vault not found") - with self.Session() as db: - new_amount = Decimal(vault.amount) + Decimal(amount) - db.query(Vault).filter_by(id=vault.id).update(amount=str(new_amount)) - db.commit() - vault = self.get_vault(wallet_id, symbol) - return vault + user = self.get_object_by_field(User, "wallet_id", wallet_id) + if not user: + raise ValueError("User not found") + return self.upsert_vault(user.id, symbol, amount) def get_vault_balance(self, wallet_id: str, symbol: str) -> str | None: """ diff --git a/quantara/web_app/db/models.py b/quantara/web_app/db/models.py index a2506f27..62b4fb7c 100644 --- a/quantara/web_app/db/models.py +++ b/quantara/web_app/db/models.py @@ -161,6 +161,10 @@ class Vault(Base): DateTime, nullable=False, default=func.now(), onupdate=func.now() ) + __table_args__ = ( + UniqueConstraint("user_id", "symbol", name="uq_vault_user_symbol"), + ) + class TransactionStatus(PyEnum): """ diff --git a/quantara/web_app/tests/db/deposit_db_tests.py b/quantara/web_app/tests/db/deposit_db_tests.py index 478142de..788acbe6 100644 --- a/quantara/web_app/tests/db/deposit_db_tests.py +++ b/quantara/web_app/tests/db/deposit_db_tests.py @@ -3,40 +3,52 @@ """ from decimal import Decimal -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch +import uuid import pytest +from sqlalchemy import create_engine, StaticPool +from sqlalchemy.orm import sessionmaker from web_app.db.crud import DepositDBConnector -from web_app.db.models import User, Vault +from web_app.db.models import Base, User, Vault @pytest.fixture -def deposit_db_connector_fixture(mock_db_connector): - """ - Fixture to provide a mocked DepositDBConnector instance using mock_db_connector. - """ - connector = DepositDBConnector() - connector.Session = mock_db_connector.Session - connector.get_object_by_field = mock_db_connector.get_object_by_field - connector.write_to_db = mock_db_connector.write_to_db +def db_session_factory(): + """Create an in-memory SQLite database for testing.""" + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + Base.metadata.create_all(engine) + Session = sessionmaker(bind=engine) + return Session + + +@pytest.fixture +def deposit_connector(db_session_factory): + """Provide a DepositDBConnector backed by an in-memory SQLite database.""" + connector = object.__new__(DepositDBConnector) + connector.Session = db_session_factory return connector @pytest.fixture -def mock_user_fixture(): +def mock_user(): """ Mocked User instance. """ - return User(id="user123", wallet_id="wallet123") + return User(id=uuid.uuid4(), wallet_id="wallet123") @pytest.fixture -def mock_vault_fixture(): +def mock_vault(): """ Mocked Vault instance. """ - return Vault(id="vault123", user_id="user123", symbol="ETH", amount="100.00") + return Vault(id=uuid.uuid4(), user_id=uuid.uuid4(), symbol="ETH", amount="100.00") class TestCreateVault: @@ -44,18 +56,11 @@ class TestCreateVault: Tests for creating a vault using DepositDBConnector. """ - def test_create_vault_success( - self, - deposit_db_connector: DepositDBConnector, - mock_user: User, - mock_db_connector, - ): + def test_create_vault_success(self, deposit_connector, mock_user): """ - Test successful creation of a vault using fixtures. + Test successful creation of a vault via upsert. """ - mock_db_connector.write_to_db = MagicMock() # Use mock_db_connector's method - - vault = deposit_db_connector.create_vault( + vault = deposit_connector.create_vault( user=mock_user, symbol="BTC", amount="50.00", @@ -64,17 +69,23 @@ def test_create_vault_success( assert vault.symbol == "BTC" assert vault.amount == "50.00" assert vault.user_id == mock_user.id - mock_db_connector.write_to_db.assert_called_once_with(vault) - def test_create_vault_failure_invalid_user( - self, - deposit_db_connector: DepositDBConnector, - ): + def test_create_vault_idempotent(self, deposit_connector, mock_user): + """ + Test that calling create_vault twice adds amounts (upsert behaviour). + """ + deposit_connector.create_vault(user=mock_user, symbol="BTC", amount="50.00") + vault = deposit_connector.create_vault( + user=mock_user, symbol="BTC", amount="25.00" + ) + assert Decimal(vault.amount) == Decimal("75.00") + + def test_create_vault_failure_invalid_user(self, deposit_connector): """ Test failure when creating a vault with an invalid user. """ - with pytest.raises(ValueError, match="Invalid user provided"): - deposit_db_connector.create_vault( + with pytest.raises((AttributeError, TypeError)): + deposit_connector.create_vault( user=None, symbol="BTC", amount="50.00", @@ -86,41 +97,24 @@ class TestAddVaultBalance: Tests for adding to a vault's balance using DepositDBConnector. """ - def test_add_balance_success( - self, - deposit_db_connector: DepositDBConnector, - mock_vault: Vault, - mock_db_connector, - ): + def test_add_balance_success(self, deposit_connector, mock_user): """ - Test successfully adding to a vault's balance using fixtures. + Test successfully adding to a vault's balance. """ - mock_db_connector.get_object_by_field = MagicMock(return_value=mock_vault) - mock_db_connector.Session().query().filter_by().update = MagicMock() - - deposit_db_connector.add_vault_balance( - wallet_id="wallet123", + deposit_connector.upsert_vault(mock_user.id, "ETH", "100.00") + vault = deposit_connector.add_vault_balance( + wallet_id=mock_user.wallet_id, symbol="ETH", amount="50.00", ) + assert Decimal(vault.amount) == Decimal("150.00") - updated_amount = Decimal(mock_vault.amount) + Decimal("50.00") - mock_db_connector.Session().query().filter_by().update.assert_called_once_with( - {"amount": str(updated_amount)} - ) - - def test_add_balance_failure_vault_not_found( - self, - deposit_db_connector: DepositDBConnector, - mock_db_connector, - ): + def test_add_balance_failure_vault_not_found(self, deposit_connector): """ - Test failure when adding to a vault balance that doesn't exist. + Test failure when adding to a vault balance for a non-existent user. """ - mock_db_connector.get_object_by_field = MagicMock(return_value=None) - - with pytest.raises(ValueError, match="Vault not found"): - deposit_db_connector.add_vault_balance( + with pytest.raises(ValueError, match="User not found"): + deposit_connector.add_vault_balance( wallet_id="invalid_wallet", symbol="ETH", amount="50.00", diff --git a/quantara/web_app/tests/test_vault.py b/quantara/web_app/tests/test_vault.py index ad989e5d..c03cb338 100644 --- a/quantara/web_app/tests/test_vault.py +++ b/quantara/web_app/tests/test_vault.py @@ -9,21 +9,9 @@ import pytest from fastapi.testclient import TestClient -from httpx import ASGITransport, AsyncClient -from web_app.api.main import app from web_app.db.crud import UserDBConnector -client = TestClient(app) - - -@pytest.fixture -async def async_client(): - """Fixture that provides an async client for testing.""" - transport = ASGITransport(app=app) - async with AsyncClient(transport=transport, base_url="http://test") as ac: - yield ac - @pytest.mark.anyio @pytest.mark.parametrize( @@ -55,7 +43,7 @@ async def test_deposit_to_vault( expected_status, expected_response, mock_user_db_connector, - async_client, + client: TestClient, ): """Test vault deposit with different scenarios.""" mock_user = MagicMock() @@ -73,9 +61,9 @@ async def test_deposit_to_vault( "web_app.db.crud.DepositDBConnector.create_vault", return_value=mock_vault, ): - response = await async_client.post("/api/vault/deposit", json=test_data) + response = client.post("/api/vault/deposit", json=test_data) else: - response = await async_client.post("/api/vault/deposit", json=test_data) + response = client.post("/api/vault/deposit", json=test_data) assert response.status_code == expected_status expected = ( @@ -112,7 +100,7 @@ async def test_get_vault_balance( balance, expected_status, expected_response, - async_client, + client: TestClient, ): """Test vault balance retrieval with different scenarios.""" with patch( @@ -120,7 +108,7 @@ async def test_get_vault_balance( return_value=balance, ): url = f"/api/vault/balance?wallet_id={wallet_id}&symbol={symbol}" - response = await async_client.get(url) + response = client.get(url) assert response.status_code == expected_status expected = ( @@ -156,7 +144,7 @@ async def test_get_vault_balance( ], ) async def test_add_vault_balance( - test_data, expected_status, expected_response, async_client + test_data, expected_status, expected_response, client: TestClient ): """Test adding to vault balance with different scenarios.""" mock_vault = MagicMock() @@ -171,7 +159,7 @@ async def test_add_vault_balance( "web_app.db.crud.DepositDBConnector.add_vault_balance", **patch_kwargs, ): - response = await async_client.post("/api/vault/add_balance", json=test_data) + response = client.post("/api/vault/add_balance", json=test_data) assert response.status_code == expected_status expected = ( diff --git a/quantara/web_app/tests/test_vault_balance.py b/quantara/web_app/tests/test_vault_balance.py new file mode 100644 index 00000000..4c211689 --- /dev/null +++ b/quantara/web_app/tests/test_vault_balance.py @@ -0,0 +1,158 @@ +""" +test_vault_balance.py +Tests for atomic vault upsert and balance update operations. +Verifies that concurrent deposits, upsert creation, upsert updates, +and unique constraint enforcement all behave correctly. +""" + +import uuid +from decimal import Decimal +from unittest.mock import MagicMock, patch, PropertyMock + +import pytest +from sqlalchemy import create_engine, StaticPool +from sqlalchemy.orm import sessionmaker + +from web_app.db.crud.deposit import DepositDBConnector +from web_app.db.models import Base, User, Vault + + +@pytest.fixture +def db_session_factory(): + """Create an in-memory SQLite database for testing.""" + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + Base.metadata.create_all(engine) + Session = sessionmaker(bind=engine) + return Session + + +@pytest.fixture +def mock_connector(db_session_factory): + """Provide a DepositDBConnector backed by an in-memory SQLite database.""" + connector = object.__new__(DepositDBConnector) + connector.Session = db_session_factory + return connector + + +@pytest.fixture +def sample_user(db_session_factory): + """Insert and return a sample User row.""" + user_id = uuid.uuid4() + user = User(wallet_id="wallet_abc", id=user_id) + with db_session_factory() as db: + db.add(user) + db.commit() + return user_id + + +class TestUpsertCreatesNewRow: + """Upssert inserts a new Vault row when none exists for the (user, symbol).""" + + def test_creates_vault_when_none_exists(self, mock_connector, sample_user): + vault = mock_connector.upsert_vault(sample_user, "ETH", "10.0") + assert vault is not None + assert vault.user_id == sample_user + assert vault.symbol == "ETH" + assert vault.amount == "10.0" + + def test_creates_multiple_vaults_for_different_symbols( + self, mock_connector, sample_user + ): + v1 = mock_connector.upsert_vault(sample_user, "ETH", "5.0") + v2 = mock_connector.upsert_vault(sample_user, "BTC", "2.0") + assert v1.symbol == "ETH" + assert v2.symbol == "BTC" + assert v1.id != v2.id + + +class TestUpsertUpdatesExistingRow: + """Upsert adds to the existing amount when a vault row already exists.""" + + def test_adds_to_existing_balance(self, mock_connector, sample_user): + mock_connector.upsert_vault(sample_user, "ETH", "10.0") + vault = mock_connector.upsert_vault(sample_user, "ETH", "5.5") + assert vault.amount == "15.5" + + def test_multiple_increments_produce_correct_sum(self, mock_connector, sample_user): + mock_connector.upsert_vault(sample_user, "XLM", "1.0") + mock_connector.upsert_vault(sample_user, "XLM", "2.0") + mock_connector.upsert_vault(sample_user, "XLM", "3.0") + vault = mock_connector.upsert_vault(sample_user, "XLM", "4.0") + assert Decimal(vault.amount) == Decimal("10.0") + + +class TestAddVaultBalance: + """add_vault_balance delegates to upsert_vault and raises on missing user.""" + + def test_adds_balance_existing_vault(self, mock_connector, sample_user): + mock_connector.upsert_vault(sample_user, "ETH", "10.0") + vault = mock_connector.add_vault_balance("wallet_abc", "ETH", "3.0") + assert Decimal(vault.amount) == Decimal("13.0") + + def test_raises_when_user_not_found(self, mock_connector): + with pytest.raises(ValueError, match="User not found"): + mock_connector.add_vault_balance("nonexistent_wallet", "ETH", "1.0") + + +class TestCreateVault: + """create_vault now delegates to upsert_vault (idempotent).""" + + def test_first_deposit(self, mock_connector, sample_user): + fake_user = MagicMock(spec=User) + fake_user.id = sample_user + vault = mock_connector.create_vault(fake_user, "ETH", "7.5") + assert Decimal(vault.amount) == Decimal("7.5") + + def test_second_deposit_adds_to_existing(self, mock_connector, sample_user): + fake_user = MagicMock(spec=User) + fake_user.id = sample_user + mock_connector.create_vault(fake_user, "ETH", "7.5") + vault = mock_connector.create_vault(fake_user, "ETH", "2.5") + assert Decimal(vault.amount) == Decimal("10.0") + + +class TestConcurrentBalanceUpdates: + """ + Simulate concurrent balance increments using SQLite transactions. + SQLite's serialized writes provide a deterministic baseline to verify + that the upsert produces the correct final sum. + """ + + def test_sequential_increments_produce_correct_final_balance( + self, mock_connector, sample_user + ): + """Simulates N sequential deposits that must all be reflected.""" + initial = Decimal("0.0") + for _ in range(10): + mock_connector.upsert_vault(sample_user, "ETH", "1.0") + initial += Decimal("1.0") + vault = mock_connector.get_vault("wallet_abc", "ETH") + assert vault is not None + assert Decimal(vault.amount) == initial + + def test_fractional_increments(self, mock_connector, sample_user): + """Verify correctness with fractional amounts.""" + amounts = ["0.1", "0.2", "0.3", "0.4"] + expected = Decimal("0.0") + for amt in amounts: + mock_connector.upsert_vault(sample_user, "ETH", amt) + expected += Decimal(amt) + vault = mock_connector.get_vault("wallet_abc", "ETH") + assert Decimal(vault.amount) == expected + + +class TestUniqueConstraint: + """The unique constraint on (user_id, symbol) is enforced by the model.""" + + def test_constraint_exists_in_model(self): + table_args = Vault.__table_args__ + constraint_names = [c.name for c in table_args if hasattr(c, "name")] + assert "uq_vault_user_symbol" in constraint_names + + def test_get_vault_returns_none_for_missing(self, mock_connector): + result = mock_connector.get_vault("no_such_wallet", "ETH") + assert result is None