Skip to content

Test for post usage when combining usage items #53

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
May 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions alembic/script.py.mako
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
"""${message}
# pylint: disable=invalid-name
"""${message if message.endswith(".") else message + "."}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Expand All @@ -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"}
32 changes: 32 additions & 0 deletions alembic/versions/9f4b6184b2f8_drop_usage_pkey.py
Original file line number Diff line number Diff line change
@@ -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 ###
844 changes: 441 additions & 403 deletions poetry.lock

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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.34.2", 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]
Expand Down
34 changes: 27 additions & 7 deletions rctab/routers/accounting/usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 postgres 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

Expand Down Expand Up @@ -78,18 +81,13 @@ 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,
)

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)
Expand All @@ -101,6 +99,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)
Expand Down Expand Up @@ -185,6 +204,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)

Expand Down
2 changes: 1 addition & 1 deletion tests/test_routes/api_calls.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
1 change: 1 addition & 0 deletions tests/test_routes/test_routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ async def create_subscription(
allocated_amount: if None then no row in allocations
(approved_amount, approved_to): if None then no row in approvals
(normal_cost, amortised_cost): the amount spent thus far
spent_date: the date the costs were incurred
"""
# pylint: disable=too-many-arguments, invalid-name

Expand Down
7 changes: 6 additions & 1 deletion tests/test_routes/test_send_emails.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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()
Expand All @@ -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()
Expand Down
Loading
Loading