From d84abaf634c7210d65bf3ade3a9bb84587235854 Mon Sep 17 00:00:00 2001 From: Iain-S <25081046+Iain-S@users.noreply.github.com> Date: Fri, 12 Jul 2024 12:28:47 +0100 Subject: [PATCH] Monthly upload route --- rctab/routers/accounting/usage.py | 103 +++++++++--------------------- tests/test_routes/test_usage.py | 101 +++++++++++++++++++++++++---- 2 files changed, 120 insertions(+), 84 deletions(-) diff --git a/rctab/routers/accounting/usage.py b/rctab/routers/accounting/usage.py index 05095a0..6e13e16 100644 --- a/rctab/routers/accounting/usage.py +++ b/rctab/routers/accounting/usage.py @@ -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 @@ -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: @@ -128,6 +123,10 @@ 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, @@ -135,76 +134,36 @@ async def post_monthly_usage( 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) diff --git a/tests/test_routes/test_usage.py b/tests/test_routes/test_usage.py index 28438b8..3e02303 100644 --- a/tests/test_routes/test_usage.py +++ b/tests/test_routes/test_usage.py @@ -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() @@ -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", @@ -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