Bank upgrades

database/crud/currency.py

@@ -5,6 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from database.crud import users
 from database.exceptions import currency as exceptions
 from database.models import Bank
+from database.utils.math.currency import rob_upgrade_price, interest_upgrade_price, capacity_upgrade_price
 
 
 NIGHTLY_AMOUNT = 420
@@ -41,3 +42,57 @@ async def claim_nightly(session: AsyncSession, user_id: int):
     session.add(bank)
     session.add(nightly_data)
     await session.commit()
+
+
+async def upgrade_capacity(database_session: AsyncSession, user_id: int) -> int:
+    """Upgrade capacity level"""
+    bank = await get_bank(database_session, user_id)
+    upgrade_price = capacity_upgrade_price(bank.capacity_level)
+
+    # Can't afford this upgrade
+    if upgrade_price > bank.dinks:
+        raise exceptions.NotEnoughDinks
+
+    bank.dinks -= upgrade_price
+    bank.capacity_level += 1
+
+    database_session.add(bank)
+    await database_session.commit()
+
+    return bank.capacity_level
+
+
+async def upgrade_interest(database_session: AsyncSession, user_id: int) -> int:
+    """Upgrade interest level"""
+    bank = await get_bank(database_session, user_id)
+    upgrade_price = interest_upgrade_price(bank.interest_level)
+
+    # Can't afford this upgrade
+    if upgrade_price > bank.dinks:
+        raise exceptions.NotEnoughDinks
+
+    bank.dinks -= upgrade_price
+    bank.interest_level += 1
+
+    database_session.add(bank)
+    await database_session.commit()
+
+    return bank.interest_level
+
+
+async def upgrade_rob(database_session: AsyncSession, user_id: int) -> int:
+    """Upgrade rob level"""
+    bank = await get_bank(database_session, user_id)
+    upgrade_price = rob_upgrade_price(bank.rob_level)
+
+    # Can't afford this upgrade
+    if upgrade_price > bank.dinks:
+        raise exceptions.NotEnoughDinks
+
+    bank.dinks -= upgrade_price
+    bank.rob_level += 1
+
+    database_session.add(bank)
+    await database_session.commit()
+
+    return bank.rob_level