From de7b5cd9609e2009383ae06b5fa6f30dd150541e Mon Sep 17 00:00:00 2001 From: stijndcl Date: Fri, 1 Mar 2024 14:18:58 +0100 Subject: [PATCH] Work on interest system --- database/crud/currency.py | 71 ++++++++++++++-- database/crud/users.py | 6 +- database/exceptions/currency.py | 4 + database/schemas.py | 19 ++++- database/utils/math/currency.py | 24 +++++- didier/cogs/currency.py | 84 ++++++++++++------- didier/cogs/tasks.py | 9 ++ tests/test_database/conftest.py | 9 ++ .../test_database/test_crud/test_currency.py | 29 ++++--- 9 files changed, 196 insertions(+), 59 deletions(-) diff --git a/database/crud/currency.py b/database/crud/currency.py index 571b5cc..38153ee 100644 --- a/database/crud/currency.py +++ b/database/crud/currency.py @@ -1,25 +1,29 @@ from datetime import date -from typing import Optional, Union +from typing import List, Optional, Union +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from database.crud import users from database.exceptions import currency as exceptions -from database.schemas import Bank, NightlyData +from database.schemas import Bank, BankSavings, NightlyData from database.utils.math.currency import ( capacity_upgrade_price, + interest_rate, interest_upgrade_price, rob_upgrade_price, + savings_cap, ) __all__ = [ "add_dinks", + "apply_daily_interest", "claim_nightly", "deduct_dinks", "gamble_dinks", "get_bank", "get_nightly_data", - "invest", + "save", "upgrade_capacity", "upgrade_interest", "upgrade_rob", @@ -36,15 +40,30 @@ async def get_bank(session: AsyncSession, user_id: int) -> Bank: return user.bank +async def get_savings(session: AsyncSession, user_id: int) -> BankSavings: + """Get a user's savings info""" + user = await users.get_or_add_user(session, user_id) + return user.savings + + async def get_nightly_data(session: AsyncSession, user_id: int) -> NightlyData: """Get a user's nightly info""" user = await users.get_or_add_user(session, user_id) return user.nightly_data -async def invest(session: AsyncSession, user_id: int, amount: Union[str, int], *, bank: Optional[Bank] = None) -> int: +async def save( + session: AsyncSession, + user_id: int, + amount: Union[str, int], + *, + bank: Optional[Bank] = None, + savings: Optional[BankSavings] = None +) -> int: """Invest some of your Dinks""" bank = bank or await get_bank(session, user_id) + savings = savings or await get_savings(session, user_id) + if amount == "all": amount = bank.dinks @@ -54,29 +73,44 @@ async def invest(session: AsyncSession, user_id: int, amount: Union[str, int], * # Don't allow investing more dinks than you own amount = min(bank.dinks, int(amount)) + # Don't allow exceeding the limit + limit = savings_cap(bank.capacity_level) + + if savings.saved >= limit: + raise exceptions.SavingsCapExceeded + + if savings.saved + amount > limit: + amount = max(0, limit - savings.saved) + bank.dinks -= amount - bank.invested += amount + savings.saved += amount session.add(bank) + session.add(savings) await session.commit() return amount async def withdraw(session: AsyncSession, user_id: int, amount: Union[str, int], *, bank: Optional[Bank] = None) -> int: - """Withdraw your invested Dinks""" + """Withdraw your saved Dinks""" bank = bank or await get_bank(session, user_id) + savings = await get_savings(session, user_id) + if amount == "all": - amount = bank.invested + amount = savings.saved # Don't allow withdrawing more dinks than you own - amount = min(bank.invested, int(amount)) + amount = min(savings.saved, int(amount)) bank.dinks += amount - bank.invested -= amount + savings.saved -= amount + savings.daily_minimum = min(savings.daily_minimum, savings.saved) session.add(bank) + session.add(savings) await session.commit() + return amount @@ -218,3 +252,22 @@ async def rob( session.add(robber) session.add(robbed) await session.commit() + + +async def apply_daily_interest(session: AsyncSession): + """Apply daily interest rates to all accounts with saved Dinks""" + statement = select(BankSavings) + all_savings: List[BankSavings] = list((await session.execute(statement)).scalars().all()) + + for savings_account in all_savings: + if savings_account.saved == 0: + continue + + bank = await get_bank(session, savings_account.user_id) + rate = interest_rate(bank.interest_level) + + savings_account.saved = float(savings_account.saved * rate) + savings_account.daily_minimum = savings_account.saved + session.add(savings_account) + + await session.commit() diff --git a/database/crud/users.py b/database/crud/users.py index bd4f2ad..9fef61d 100644 --- a/database/crud/users.py +++ b/database/crud/users.py @@ -3,7 +3,7 @@ from typing import Optional from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from database.schemas import Bank, NightlyData, User +from database.schemas import Bank, BankSavings, NightlyData, User __all__ = [ "get_or_add_user", @@ -37,8 +37,12 @@ async def get_or_add_user(session: AsyncSession, user_id: int, *, options: Optio user.bank = bank user.nightly_data = nightly_data + savings = BankSavings(user_id=user_id) + user.savings = savings + session.add(bank) session.add(nightly_data) + session.add(savings) session.add(user) await session.commit() diff --git a/database/exceptions/currency.py b/database/exceptions/currency.py index c06c76f..2cc9a1d 100644 --- a/database/exceptions/currency.py +++ b/database/exceptions/currency.py @@ -7,3 +7,7 @@ class DoubleNightly(Exception): class NotEnoughDinks(Exception): """Exception raised when trying to do something you don't have the Dinks for""" + + +class SavingsCapExceeded(Exception): + """Exception raised when trying to save more Dinks than the cap allows""" diff --git a/database/schemas.py b/database/schemas.py index 2dc421a..88c4174 100644 --- a/database/schemas.py +++ b/database/schemas.py @@ -12,6 +12,7 @@ from database import enums __all__ = [ "Base", "Bank", + "BankSavings", "Birthday", "Bookmark", "CommandStats", @@ -52,7 +53,6 @@ class Bank(Base): user_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("users.user_id")) dinks: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False) - invested: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False) # Interest rate interest_level: Mapped[int] = mapped_column(server_default="1", nullable=False) @@ -66,6 +66,20 @@ class Bank(Base): user: Mapped[User] = relationship(uselist=False, back_populates="bank", lazy="selectin") +class BankSavings(Base): + """Savings information for a user's bank""" + + __tablename__ = "savings" + + savings_id: Mapped[int] = mapped_column(primary_key=True) + user_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("users.user_id")) + + saved: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False) + daily_minimum: Mapped[int] = mapped_column(BigInteger, server_default="0", nullable=False) + + user: Mapped[User] = relationship(uselist=False, back_populates="savings", lazy="selectin") + + class Birthday(Base): """A user's birthday""" @@ -350,3 +364,6 @@ class User(Base): reminders: Mapped[List[Reminder]] = relationship( back_populates="user", uselist=True, lazy="selectin", cascade="all, delete-orphan" ) + savings: Mapped[List[BankSavings]] = relationship( + back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan" + ) diff --git a/database/utils/math/currency.py b/database/utils/math/currency.py index a8747a2..ff347a4 100644 --- a/database/utils/math/currency.py +++ b/database/utils/math/currency.py @@ -4,6 +4,8 @@ __all__ = [ "capacity_upgrade_price", "interest_upgrade_price", "rob_upgrade_price", + "interest_rate", + "savings_cap", "jail_chance", "jail_time", "rob_amount", @@ -16,7 +18,7 @@ def interest_upgrade_price(level: int) -> int: base_cost = 600 growth_rate = 1.8 - return math.floor(base_cost * (growth_rate**level)) + return math.floor(base_cost * (growth_rate ** (level - 1))) def capacity_upgrade_price(level: int) -> int: @@ -24,7 +26,7 @@ def capacity_upgrade_price(level: int) -> int: base_cost = 800 growth_rate = 1.6 - return math.floor(base_cost * (growth_rate**level)) + return math.floor(base_cost * (growth_rate ** (level - 1))) def rob_upgrade_price(level: int) -> int: @@ -32,7 +34,23 @@ def rob_upgrade_price(level: int) -> int: base_cost = 950 growth_rate = 1.9 - return math.floor(base_cost * (growth_rate**level)) + return math.floor(base_cost * (growth_rate ** (level - 1))) + + +def interest_rate(level: int) -> float: + """Calculate the amount of interest you will receive""" + base_rate = 1.025 + growth_rate = 0.03 + + return base_rate + (growth_rate * (level - 1)) + + +def savings_cap(level: int) -> int: + """Calculate the maximum amount you can save""" + base_limit = 1000 + growth_rate = 1.10 + + return math.floor(base_limit * (growth_rate ** (level - 1))) def jail_chance(level: int) -> float: diff --git a/didier/cogs/currency.py b/didier/cogs/currency.py index 5828237..69e6363 100644 --- a/didier/cogs/currency.py +++ b/didier/cogs/currency.py @@ -10,16 +10,23 @@ from discord.ext import commands import settings from database.crud import currency as crud +from database.crud import users from database.crud.jail import get_user_jail, imprison -from database.exceptions.currency import DoubleNightly, NotEnoughDinks +from database.exceptions.currency import ( + DoubleNightly, + NotEnoughDinks, + SavingsCapExceeded, +) from database.utils.math.currency import ( capacity_upgrade_price, + interest_rate, interest_upgrade_price, jail_chance, jail_time, rob_amount, rob_chance, rob_upgrade_price, + savings_cap, ) from didier import Didier from didier.utils.discord import colours @@ -51,17 +58,20 @@ class Currency(commands.Cog): """Award a user `amount` Didier Dinks.""" async with self.client.postgres_session as session: await crud.add_dinks(session, user.id, amount) - plural = pluralize("Didier Dink", amount) - await ctx.reply( - f"{ctx.author.display_name} has awarded **{user.display_name}** with **{amount}** {plural}.", - mention_author=False, - ) + + plural = pluralize("Didier Dink", amount) + await ctx.reply( + f"{ctx.author.display_name} has awarded **{user.display_name}** with **{amount}** {plural}.", + mention_author=False, + ) @commands.group(name="bank", aliases=["b"], case_insensitive=True, invoke_without_command=True) async def bank(self, ctx: commands.Context): """Show your Didier Bank information.""" async with self.client.postgres_session as session: - bank = await crud.get_bank(session, ctx.author.id) + user = await users.get_or_add_user(session, ctx.author.id) + bank = user.bank + savings = user.savings embed = discord.Embed(title="Bank of Didier", colour=discord.Colour.blue()) embed.set_author(name=ctx.author.display_name) @@ -69,9 +79,10 @@ class Currency(commands.Cog): if ctx.author.avatar is not None: embed.set_thumbnail(url=ctx.author.avatar.url) - embed.add_field(name="Interest level", value=bank.interest_level) - embed.add_field(name="Capacity level", value=bank.capacity_level) - embed.add_field(name="Currently invested", value=bank.invested, inline=False) + embed.add_field(name="Interest rate", value=round(interest_rate(bank.interest_level), 2)) + embed.add_field(name="Maximum capacity", value=round(savings_cap(bank.capacity_level), 2)) + embed.add_field(name="Currently saved", value=savings.saved, inline=False) + embed.add_field(name="Daily minimum", value=savings.daily_minimum, inline=False) await ctx.reply(embed=embed, mention_author=False) @@ -135,52 +146,61 @@ class Currency(commands.Cog): """Check your Didier Dinks.""" async with ctx.typing(), self.client.postgres_session as session: bank = await crud.get_bank(session, ctx.author.id) - plural = pluralize("Didier Dink", bank.dinks) - await ctx.reply(f"You have **{bank.dinks}** {plural}.", mention_author=False) - @commands.command(name="invest", aliases=["deposit", "dep", "i"]) # type: ignore[arg-type] - async def invest(self, ctx: commands.Context, amount: typing.Annotated[typing.Union[str, int], abbreviated_number]): - """Invest `amount` Didier Dinks into your bank. + plural = pluralize("Didier Dink", bank.dinks) + await ctx.reply(f"You have **{bank.dinks}** {plural}.", mention_author=False) + + @commands.command(name="save", aliases=["deposit", "dep", "s"]) # type: ignore[arg-type] + async def save(self, ctx: commands.Context, amount: typing.Annotated[typing.Union[str, int], abbreviated_number]): + """Add `amount` Didier Dinks into your bank's savings account. The `amount`-argument can take both raw numbers, and abbreviations of big numbers. Additionally, passing `all` or `*` as the value will invest all of your Didier Dinks. Example usage: ``` - didier invest all - didier invest 500 - didier invest 25k - didier invest 5.3b + didier save all + didier save 500 + didier save 25k + didier save 5.3b ``` """ if isinstance(amount, int) and amount <= 0: return await ctx.reply("Amount of Didier Dinks to invest must be a strictly positive integer.") async with self.client.postgres_session as session: - invested = await crud.invest(session, ctx.author.id, amount) - plural = pluralize("Didier Dink", invested) + try: + saved = await crud.save(session, ctx.author.id, amount) + except SavingsCapExceeded: + return await ctx.reply( + "You have already exceeded the savings cap for your level. Upgrade your bank's capacity to save " + "more." + ) - if invested == 0: - await ctx.reply("You don't have any Didier Dinks to invest.", mention_author=False) - else: - await ctx.reply(f"You have invested **{invested}** {plural}.", mention_author=False) + plural = pluralize("Didier Dink", saved) - @commands.command(name="withdraw", aliases=["uninvest", "w"]) # type: ignore[arg-type] + if saved == 0: + await ctx.reply("You don't have any Didier Dinks to invest.", mention_author=False) + else: + await ctx.reply(f"You have saved **{saved}** {plural}.", mention_author=False) + + @commands.command(name="withdraw", aliases=["undeposit", "unsave", "w"]) # type: ignore[arg-type] async def withdraw( self, ctx: commands.Context, amount: typing.Annotated[typing.Union[str, int], abbreviated_number] ): - """Withdraw some of your invested Didier Dinks from your bank.""" + """Withdraw some of your Didier Dinks from your bank's savings account.""" if isinstance(amount, int) and amount <= 0: return await ctx.reply("Amount of Didier Dinks to invest must be a strictly positive integer.") async with self.client.postgres_session as session: withdrawn = await crud.withdraw(session, ctx.author.id, amount) - plural = pluralize("Didier Dink", withdrawn) - if withdrawn == 0: - await ctx.reply("You don't have any Didier Dinks to withdraw.", mention_author=False) - else: - await ctx.reply(f"You have withdrawn **{withdrawn}** {plural}.", mention_author=False) + plural = pluralize("Didier Dink", withdrawn) + + if withdrawn == 0: + await ctx.reply("You don't have any Didier Dinks to withdraw.", mention_author=False) + else: + await ctx.reply(f"You have withdrawn **{withdrawn}** {plural}.", mention_author=False) @commands.hybrid_command(name="nightly") # type: ignore[arg-type] async def nightly(self, ctx: commands.Context): diff --git a/didier/cogs/tasks.py b/didier/cogs/tasks.py index f59d697..778f47e 100644 --- a/didier/cogs/tasks.py +++ b/didier/cogs/tasks.py @@ -10,6 +10,7 @@ from overrides import overrides import settings from database import enums from database.crud.birthdays import get_birthdays_on_day +from database.crud.currency import apply_daily_interest from database.crud.reminders import get_all_reminders_for_category from database.crud.ufora_announcements import remove_old_announcements from database.schemas import Reminder @@ -77,6 +78,7 @@ class Tasks(commands.Cog): # Start other tasks self.reminders.start() + self.daily_interest.start() asyncio.create_task(self.get_error_channel()) @overrides @@ -318,11 +320,18 @@ class Tasks(commands.Cog): async with self.client.postgres_session as session: await remove_old_announcements(session) + @tasks.loop(time=DAILY_RESET_TIME) + async def daily_interest(self): + """Give everyone's daily interest""" + async with self.client.postgres_session as session: + await apply_daily_interest(session) + @check_birthdays.error @pull_schedules.error @pull_ufora_announcements.error @reminders.error @remove_old_ufora_announcements.error + @daily_interest.error async def _on_tasks_error(self, error: BaseException): """Error handler for all tasks""" self.client.dispatch("task_error", error) diff --git a/tests/test_database/conftest.py b/tests/test_database/conftest.py index a675a35..6cf41d4 100644 --- a/tests/test_database/conftest.py +++ b/tests/test_database/conftest.py @@ -6,6 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from database.crud import users from database.schemas import ( Bank, + BankSavings, UforaAnnouncement, UforaCourse, UforaCourseAlias, @@ -38,6 +39,14 @@ async def bank(postgres: AsyncSession, user: User) -> Bank: return _bank +@pytest.fixture +async def savings(postgres: AsyncSession, user: User) -> BankSavings: + """Fixture to fetch the test user's savings account""" + _savings = user.savings + await postgres.refresh(_savings) + return _savings + + @pytest.fixture async def ufora_course(postgres: AsyncSession) -> UforaCourse: """Fixture to create a course""" diff --git a/tests/test_database/test_crud/test_currency.py b/tests/test_database/test_crud/test_currency.py index 1beddc6..6cbb8d4 100644 --- a/tests/test_database/test_crud/test_currency.py +++ b/tests/test_database/test_crud/test_currency.py @@ -6,7 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from database.crud import currency as crud from database.exceptions import currency as exceptions -from database.schemas import Bank +from database.schemas import Bank, BankSavings async def test_add_dinks(postgres: AsyncSession, bank: Bank): @@ -40,40 +40,43 @@ async def test_claim_nightly_unavailable(postgres: AsyncSession, bank: Bank): assert bank.dinks == crud.NIGHTLY_AMOUNT -async def test_invest(postgres: AsyncSession, bank: Bank): - """Test investing some Dinks""" +async def test_save(postgres: AsyncSession, bank: Bank, savings: BankSavings): + """Test saving some Dinks""" bank.dinks = 100 postgres.add(bank) await postgres.commit() - await crud.invest(postgres, bank.user_id, 20) + await crud.save(postgres, bank.user_id, 20, bank=bank, savings=savings) await postgres.refresh(bank) + await postgres.refresh(savings) assert bank.dinks == 80 - assert bank.invested == 20 + assert savings.saved == 20 -async def test_invest_all(postgres: AsyncSession, bank: Bank): - """Test investing all dinks""" +async def test_save_all(postgres: AsyncSession, bank: Bank, savings: BankSavings): + """Test saving all dinks""" bank.dinks = 100 postgres.add(bank) await postgres.commit() - await crud.invest(postgres, bank.user_id, "all") + await crud.save(postgres, bank.user_id, "all", bank=bank, savings=savings) await postgres.refresh(bank) + await postgres.refresh(savings) assert bank.dinks == 0 - assert bank.invested == 100 + assert savings.saved == 100 -async def test_invest_more_than_owned(postgres: AsyncSession, bank: Bank): - """Test investing more Dinks than you own""" +async def test_save_more_than_owned(postgres: AsyncSession, bank: Bank, savings: BankSavings): + """Test saving more Dinks than you own""" bank.dinks = 100 postgres.add(bank) await postgres.commit() - await crud.invest(postgres, bank.user_id, 200) + await crud.save(postgres, bank.user_id, 200, bank=bank, savings=savings) await postgres.refresh(bank) + await postgres.refresh(savings) assert bank.dinks == 0 - assert bank.invested == 100 + assert savings.saved == 100