Test for post usage when combining usage items

alembic/script.py.mako

@@ -1,4 +1,5 @@
-"""${message}
+# pylint: disable=invalid-name
+"""${message if message.endswith(".") else message + "."}
 
 Revision ID: ${up_revision}
 Revises: ${down_revision | comma,n}
@@ -9,16 +10,20 @@ from alembic import op
 import sqlalchemy as sa
 ${imports if imports else ""}
 
+# pylint: disable=no-member
+
 # revision identifiers, used by Alembic.
 revision = ${repr(up_revision)}
 down_revision = ${repr(down_revision)}
 branch_labels = ${repr(branch_labels)}
 depends_on = ${repr(depends_on)}
 
 
-def upgrade():
+def upgrade() -> None:
+    """Upgrade the database."""
     ${upgrades if upgrades else "pass"}
 
 
-def downgrade():
+def downgrade() -> None:
+    """Downgrade the database."""
     ${downgrades if downgrades else "pass"}

alembic/versions/9f4b6184b2f8_drop_usage_pkey.py

@@ -0,0 +1,32 @@
+# pylint: disable=invalid-name
+"""Drop the primary key of the usage table.
+
+Revision ID: 9f4b6184b2f8
+Revises: b65796c99771
+Create Date: 2024-10-10 11:28:39.846967
+
+"""
+
+from alembic import op
+
+# pylint: disable=no-member
+
+# revision identifiers, used by Alembic.
+revision = "9f4b6184b2f8"
+down_revision = "b65796c99771"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade the database."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint("usage_pkey", "usage", type_="primary", schema="accounting")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade the database."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_primary_key("usage_pkey", "usage", ["id"], schema="accounting")
+    # ### end Alembic commands ###

pyproject.toml

@@ -25,7 +25,6 @@ pydantic = { extras = ["email"], version = "^2.7.1" }
 pydantic-settings = "^2.3.4"
 PyJWT = { extras = ["crypto"], version = "^2.4.0" }
 python-dotenv = "^1.0.1"
-rctab_models = { git = "https://github.com/alan-turing-institute/rctab-models", tag = "0.1.0" }
 redis = {extras = ["hiredis"], version = "^5.0.1"}
 requests = "^2.32.3"
 sendgrid = "^6.9.1"
@@ -35,6 +34,7 @@ sphinx-rtd-theme = {version = "^1.3.0", optional = true}
 sphinxcontrib-napoleon = {version = "^0.7", optional = true}
 secure = "^0.3.0"
 uvicorn = { version = "^0.30.1", extras = ["standard"] }
+rctab-models = { git = "https://github.com/alan-turing-institute/rctab-models", tag = "0.2.0" }
 fastapi = "^0.115.6"
 
 [tool.poetry.group.dev.dependencies]

rctab/routers/accounting/usage.py

@@ -10,7 +10,10 @@
 from fastapi.security import OAuth2PasswordBearer
 from pydantic import BaseModel
 from rctab_models.models import AllCMUsage, AllUsage, CMUsage, Usage, UserRBAC
-from sqlalchemy import select
+from sqlalchemy import delete, select
+
+# require the postgre specific insert rather than the generic sqlachemy for;
+# `post_cm_usage` fn where "excluded" and "on_conflict_do_update" are used.
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.inspection import inspect
 
@@ -78,18 +81,18 @@ async def insert_usage(all_usage: AllUsage) -> None:
         all_usage: Usage data to insert.
     """
     usage_query = insert(accounting_models.usage)
-    update_dict = {c.name: c for c in usage_query.excluded if not c.primary_key}
-    on_duplicate_key_stmt = usage_query.on_conflict_do_update(
-        index_elements=inspect(accounting_models.usage).primary_key,
-        set_=update_dict,
-    )
+    # update_dict = {c.name: c for c in usage_query.excluded if not c.primary_key}
+    # on_duplicate_key_stmt = usage_query.on_conflict_do_update(
+    #     index_elements=inspect(accounting_models.usage).primary_key,
+    #     set_=update_dict,
+    # )
 
     logger.info("Inserting usage data")
     insert_start = datetime.datetime.now()
 
     await executemany(
         database,
-        on_duplicate_key_stmt,
+        usage_query,
         values=[i.model_dump() for i in all_usage.usage_list],
     )
     logger.info("Inserting usage data took %s", datetime.datetime.now() - insert_start)
@@ -101,6 +104,27 @@ async def insert_usage(all_usage: AllUsage) -> None:
     )
 
 
+async def delete_usage(start_date: datetime.date, end_date: datetime.date) -> None:
+    """Deletes usage(s) within a date range from the database.
+
+    Args:
+        start_date: Defines the beginning of the date range.
+        end_date: Defines the end of the date range.
+    """
+    usage_query = (
+        delete(accounting_models.usage)
+        .where(accounting_models.usage.c.date >= start_date)
+        .where(accounting_models.usage.c.date <= end_date)
+    )
+
+    logger.info("Delete usage data within a date range")
+    delete_start = datetime.datetime.now()
+
+    await database.execute(usage_query)
+
+    logger.info("Delete usage data took %s", datetime.datetime.now() - delete_start)
+
+
 @router.post("/monthly-usage", response_model=TmpReturnStatus)
 async def post_monthly_usage(
     all_usage: AllUsage, _: Dict[str, str] = Depends(authenticate_usage_app)
@@ -185,6 +209,7 @@ async def post_usage(
             unique_subscriptions = list(
                 {i.subscription_id for i in all_usage.usage_list}
             )
+            await delete_usage(all_usage.start_date, all_usage.end_date)
 
             await insert_subscriptions_if_not_exists(unique_subscriptions)
 
@@ -199,6 +224,8 @@ async def post_usage(
     )
 
 
+# TODO: remove this decorator and leave it as a standard function as it's only
+# used in test. Hence, it will prevent any accidental call to the function.
 @router.get("/all-usage", response_model=List[Usage])
 async def get_usage(_: UserRBAC = Depends(token_admin_verified)) -> List[Usage]:
     """Get all usage data."""

tests/test_routes/api_calls.py

@@ -173,7 +173,7 @@ def create_usage(
         monthly_upload=None,
     )
 
-    post_data = AllUsage(usage_list=[usage])
+    post_data = AllUsage(usage_list=[usage], start_date=date, end_date=date)
 
     return client.post(
         "usage/all-usage",

tests/test_routes/test_send_emails.py

@@ -126,6 +126,7 @@ async def test_usage_emails(
         allocated_amount=100.0,
         approved=(100.0, date.today() + timedelta(days=10)),
         spent=(20.0, 0),
+        spent_date=date.today() - timedelta(days=5),
     )
 
     ninety_percent = await create_subscription(
@@ -135,6 +136,7 @@ async def test_usage_emails(
         allocated_amount=100.0,
         approved=(100.0, date.today() + timedelta(days=10)),
         spent=(80.0, 0),
+        spent_date=date.today() - timedelta(days=5),
     )
 
     ninety_five_percent = await create_subscription(
@@ -144,6 +146,7 @@ async def test_usage_emails(
         allocated_amount=100.0,
         approved=(100.0, date.today() + timedelta(days=10)),
         spent=(85.0, 0),
+        spent_date=date.today() - timedelta(days=5),
     )
 
     ninety_percent_usage = USAGE_DICT.copy()
@@ -169,7 +172,9 @@ async def test_usage_emails(
             Usage(**ninety_percent_usage),
             Usage(**thirty_percent_usage),
             Usage(**ninety_five_percent_usage),
-        ]
+        ],
+        start_date=date.today() - timedelta(3),
+        end_date=date.today() - timedelta(2),
     )
 
     mock_send = AsyncMock()