Skip to content

Commit

Permalink
Monthly upload route
Browse files Browse the repository at this point in the history
  • Loading branch information
Iain-S committed Jul 12, 2024
1 parent c921f7d commit d84abaf
Show file tree
Hide file tree
Showing 2 changed files with 120 additions and 84 deletions.
103 changes: 31 additions & 72 deletions rctab/routers/accounting/usage.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
"""Set and get usage data."""

import calendar
import datetime
import logging
from functools import reduce
from typing import Dict, List
from uuid import UUID

Expand Down Expand Up @@ -110,16 +108,13 @@ async def post_monthly_usage(
"""Inserts monthly usage data into the database."""
logger.info("Post monthly usage called")

post_start = datetime.datetime.now()
if len(all_usage.usage_list) == 0:
raise HTTPException(
status_code=400,
detail="Monthly usage data must have at least one record.",
)

dates = sorted([x.date for x in all_usage.usage_list])
# if len(dates) == 0:
# raise HTTPException(
# status_code=400,
# detail="Post monthly usage data must have at least one record.",
# )
date_min = dates[-1]
date_max = dates[-0]
post_start = datetime.datetime.now()

for usage in all_usage.usage_list:
if usage.monthly_upload is None:
Expand All @@ -128,83 +123,47 @@ async def post_monthly_usage(
detail="Post monthly usage data must have the monthly_upload column populated.",
)

dates = sorted([x.date for x in all_usage.usage_list])
date_min = dates[0]
date_max = dates[-1]

logger.info(
"Post monthly usage received data for %s - %s containing %d records",
date_min,
date_max,
len(all_usage.usage_list),
)

if date_min.year != date_max.year or date_min.month != date_max.month:
raise HTTPException(
status_code=400,
detail=(
"Post monthly usage data should contain usage only for one month. "
f"Min, Max usage date: ({str(date_min)}), ({str(date_max)})."
),
)

month_start = datetime.date(date_min.year, date_min.month, 1)
month_end = datetime.date(
date_min.year,
date_min.month,
calendar.monthrange(date_min.year, date_min.month)[1],
)

logger.info(
"Post monthly usage checks if data for %s - %s has already been posted",
month_start,
month_end,
)

# Check if monthly usage has already been posted for the month
query = select([accounting_models.usage])
query = query.where(accounting_models.usage.c.date >= month_start)
query = query.where(accounting_models.usage.c.date <= month_end)
query = query.where(accounting_models.usage.c.monthly_upload.isnot(None))

query_result = await database.fetch_all(query)
async with database.transaction():

if query_result is not None and len(query_result) > 0:
raise HTTPException(
status_code=400,
detail=f"Post monthly usage data for {str(month_start)}-{str(month_end)} has already been posted.",
logger.info(
"Post monthly usage deleting existing usage data for %s - %s",
date_min,
date_max,
)

async with UsageEmailContextManager(database):

async with database.transaction():

logger.info(
"Post monthly usage deleting existing usage data for %s - %s",
month_start,
month_end,
)

# delete al the usage for the month
query_del = accounting_models.usage.delete().where(
accounting_models.usage.c.date >= month_start
)
query_del = query_del.where(accounting_models.usage.c.date <= month_end)
await database.execute(query_del)

logger.info(
"Post monthly usage inserting new subscriptions if they don't exist"
)
# Delete all usage for the time period to have a blank slate.
query_del = (
accounting_models.usage.delete()
.where(accounting_models.usage.c.date >= date_min)
.where(accounting_models.usage.c.date <= date_max)
)
await database.execute(query_del)

unique_subscriptions = list(
{i.subscription_id for i in all_usage.usage_list}
)
logger.info(
"Post monthly usage inserting new subscriptions if they don't exist"
)

await insert_subscriptions_if_not_exists(unique_subscriptions)
unique_subscriptions = list({i.subscription_id for i in all_usage.usage_list})

logger.info("Post monthly usage inserting monthly usage data")
await insert_subscriptions_if_not_exists(unique_subscriptions)

await insert_usage(all_usage)
logger.info("Post monthly usage inserting monthly usage data")

logger.info("Post monthly usage refreshing desired states")
await insert_usage(all_usage)

await refresh_desired_states(UUID(ADMIN_OID), unique_subscriptions)
# Note that we don't refresh the desired states here as we don't
# want to trigger excess emails.

logger.info("Post monthly usage data took %s", datetime.datetime.now() - post_start)

Expand Down
101 changes: 89 additions & 12 deletions tests/test_routes/test_usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,14 @@
BillingStatus,
CMUsage,
SubscriptionDetails,
Usage,
)
from rctab.routers.accounting.usage import post_usage
from rctab.routers.accounting.usage import get_usage, post_monthly_usage, post_usage
from tests.test_routes import api_calls, constants
from tests.test_routes.test_routes import test_db # pylint: disable=unused-import
from tests.test_routes.test_routes import ( # pylint: disable=unused-import
create_subscription,
test_db,
)
from tests.utils import print_list_diff

date_from = datetime.date.today()
Expand Down Expand Up @@ -301,12 +305,12 @@ def test_post_monthly_usage(
# Should error if there is no data.
resp = client.post(
"usage/monthly-usage",
content=AllUsage.model_dump_json().encode("utf-8"),
content=AllUsage(usage_list=[]).model_dump_json().encode("utf-8"),
headers={"authorization": "Bearer " + token},
)

assert resp.status_code == 400
assert "usage list is empty" in resp.text
assert "must have at least one record" in resp.text

resp = client.post(
"usage/monthly-usage",
Expand Down Expand Up @@ -356,17 +360,90 @@ def test_post_monthly_usage(


@pytest.mark.asyncio
async def test_monthly_usage_dates(
test_db: Database, mocker: MockerFixture # pylint: disable=redefined-outer-name
async def test_monthly_usage_2(
test_db: Database, # pylint: disable=redefined-outer-name
) -> None:
mock_refresh = AsyncMock()
mocker.patch(
"rctab.routers.accounting.usage.refresh_materialised_view", mock_refresh
sub1 = await create_subscription(test_db)
sub2 = await create_subscription(test_db)

await post_usage(
AllUsage(
usage_list=[
Usage(
id=str(UUID(int=0)),
subscription_id=sub1,
date="2024-04-01",
total_cost=1.0,
invoice_section="-",
),
Usage(
id=str(UUID(int=1)),
subscription_id=sub2,
date="2024-04-02",
total_cost=2.0,
invoice_section="-",
),
Usage(
id=str(UUID(int=2)),
subscription_id=sub1,
date="2024-04-03",
total_cost=4.0,
invoice_section="-",
),
]
),
{"mock": "authentication"},
)

await post_usage(AllUsage(usage_list=[]), {"mock": "authentication"})

mock_refresh.assert_called_once_with(test_db, usage_view)
await post_monthly_usage(
AllUsage(
usage_list=[
Usage(
id=str(UUID(int=3)),
subscription_id=sub1,
date="2024-04-01",
total_cost=10.0,
invoice_section="-",
monthly_upload=datetime.date.today(),
),
Usage(
id=str(UUID(int=4)),
subscription_id=sub1,
date="2024-04-02",
total_cost=0.5,
invoice_section="-",
monthly_upload=datetime.date.today(),
),
]
),
{"mock": "authentication"},
)
all_usage = await get_usage()
assert all_usage == [
Usage(
id=str(UUID(int=2)),
subscription_id=sub1,
date="2024-04-03",
total_cost=4.0,
invoice_section="-",
),
Usage(
id=str(UUID(int=3)),
subscription_id=sub1,
date="2024-04-01",
total_cost=10.0,
invoice_section="-",
monthly_upload=datetime.date.today(),
),
Usage(
id=str(UUID(int=4)),
subscription_id=sub1,
date="2024-04-02",
total_cost=0.5,
invoice_section="-",
monthly_upload=datetime.date.today(),
),
]


@pytest.mark.asyncio
Expand Down

0 comments on commit d84abaf

Please sign in to comment.