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

database/exceptions/currency.py

@@ -1,2 +1,6 @@
 class DoubleNightly(Exception):
     """Exception raised when claiming nightlies multiple times per day"""
+
+
+class NotEnoughDinks(Exception):
+    """Exception raised when trying to do something you don't have the Dinks for"""

database/utils/math/currency.py

@@ -0,0 +1,25 @@
+import math
+
+
+def interest_upgrade_price(level: int) -> int:
+    """Calculate the price to upgrade your interest level"""
+    base_cost = 600
+    growth_rate = 1.8
+
+    return math.floor(base_cost * (growth_rate**level))
+
+
+def capacity_upgrade_price(level: int) -> int:
+    """Calculate the price to upgrade your capacity level"""
+    base_cost = 800
+    growth_rate = 1.6
+
+    return math.floor(base_cost * (growth_rate**level))
+
+
+def rob_upgrade_price(level: int) -> int:
+    """Calculate the price to upgrade your rob level"""
+    base_cost = 950
+    growth_rate = 1.9
+
+    return math.floor(base_cost * (growth_rate**level))