[MPT-10795] Add TrialInfo dataclass and code improvements

Author: antoniodimariano

ffc/billing/dataclasses.py

@@ -25,6 +25,13 @@ class Refund:
     description: str
 
 
+@dataclass
+class TrialInfo:
+    trial_days: set[int] | None = None
+    refund_from: date | None = None
+    refund_to: date | None = None
+
+
 @dataclass
 class CurrencyConversionInfo:
     base_currency: str

ffc/process_billing.py ffc/billing/process_billing.pyffc/process_billing.py renamed to ffc/billing/process_billing.py

@@ -7,7 +7,7 @@
 from datetime import UTC, date, datetime
 from io import BytesIO
 from pathlib import Path
-from typing import Any
+from typing import Any, AsyncGenerator, Coroutine
 
 import aiofiles
 import aiofiles.os
@@ -22,6 +22,7 @@
     CurrencyConversionInfo,
     Datasource,
     Refund,
+    TrialInfo,
 )
 from ffc.billing.exceptions import ExchangeRatesClientError, JournalStatusError
 from ffc.clients.exchage_rates import ExchangeRatesAsyncClient
@@ -112,7 +113,7 @@ def __init__(
         )
 
     @asynccontextmanager
-    async def acquire_semaphore(self):
+    async def acquire_semaphore(self) -> AsyncGenerator[None, Any]:
         """
         This method acquires and releases a  semaphore.
         """
@@ -129,7 +130,7 @@ async def maybe_call(
         func,
         *args,
         **kwargs,
-    ):
+    ) -> Any | None:
         """
         Conditionally calls and awaits the given asynchronous function.
 
@@ -149,7 +150,7 @@ async def maybe_call(
         if not self.dry_run:
             return await func(*args, **kwargs)
 
-    def build_filepath(self):
+    def build_filepath(self) -> str:
         """
         Constructs and returns the file path for a charges JSONL file.
 
@@ -166,7 +167,7 @@ def build_filepath(self):
         filepath = f"{tempfile.gettempdir()}/{filepath}" if not self.dry_run else filepath
         return filepath
 
-    async def evaluate_journal_status(self, journal_external_id):
+    async def evaluate_journal_status(self, journal_external_id) -> dict[str, Any] | None:
         """
         Evaluates the status of a journal.
 
@@ -195,7 +196,7 @@ async def evaluate_journal_status(self, journal_external_id):
             self.logger.warning(f"Found the journal {journal_id} with status {journal_status}")
             raise JournalStatusError()
 
-    async def process(self):
+    async def process(self) -> None:
         """
         This method is responsible for passing a journal with status VALIDATED to
         the function that writes the charges into a file and then to complete the
@@ -204,7 +205,6 @@ async def process(self):
         result = AuthorizationProcessResult(authorization_id=self.authorization_id)
         async with self.acquire_semaphore():
             try:
-                # double check with production
                 if not await self.mpt_client.count_active_agreements(
                     self.authorization_id,
                     self.billing_start_date,
@@ -248,7 +248,7 @@ async def process(self):
             except Exception as error:
                 self.logger.error(f"An error occurred: {error}", exc_info=error)
 
-    async def write_charges_file(self, filepath):
+    async def write_charges_file(self, filepath) -> bool:
         """
         This method writes the charges file to the given filepath.
         If there is more than one agreement for an organization, it won't be processed.
@@ -308,7 +308,9 @@ async def write_charges_file(self, filepath):
                 return False
             return True
 
-    async def complete_journal_process(self, filepath, journal, journal_external_id):
+    async def complete_journal_process(
+        self, filepath, journal, journal_external_id
+    ) -> Coroutine | None:
         """
         This method uploads and submits the given journal, attaching also the exchange rates
         files.
@@ -377,7 +379,7 @@ async def get_currency_conversion_info(
 
     async def attach_exchange_rates(
         self, journal_id: str, currency: str, exchange_rates: dict[str, Any]
-    ):
+    ) -> Coroutine | None:
         """
         This method checks if an attachment already exists for the given journal.
         If it exists, it will be deleted and a new one will be created with the
@@ -392,7 +394,7 @@ async def attach_exchange_rates(
         attachment = await self.mpt_client.fetch_journal_attachment(journal_id, f"{currency}_")
         if attachment:  # pragma no cover
             if attachment["name"] == filename:
-                return
+                return None
             await self.mpt_client.delete_journal_attachment(journal_id, attachment["id"])
 
         return await self.mpt_client.create_journal_attachment(journal_id, filename, serialized)
@@ -402,7 +404,7 @@ async def dump_organization_charges(
         charges_file: Any,
         organization: dict[str, Any],
         agreement: dict[str, Any] | None = None,
-    ):
+    ) -> Coroutine | None:
         organization_id = organization["id"]
         async for datasource_info, expenses in async_groupby(
             self.ffc_client.fetch_organization_expenses(organization_id, self.year, self.month),
@@ -441,7 +443,7 @@ async def generate_datasource_charges(
         datasource_id: str,
         datasource_name: str,
         daily_expenses: dict[int, Decimal],
-    ):
+    ) -> list[str]:
         """
         This method generates all the charges for the given datasource and
         calculates the refund for the Trials and Entitlements periods.
@@ -539,7 +541,7 @@ async def generate_datasource_charges(
         )
         return charges
 
-    async def is_journal_status_validated(self, journal_id, max_attempts=5):
+    async def is_journal_status_validated(self, journal_id, max_attempts=5) -> bool:
         backoff_times = [0.15, 0.45, 1.05, 2.25, 4.65]
 
         for attempt in range(min(max_attempts, len(backoff_times))):
@@ -566,17 +568,19 @@ def generate_refunds(
         trial_days = set()
         trial_start_date, trial_end_date = get_trial_dates(agreement=agreement)
         if trial_start_date and trial_end_date:
-            trial_days, trial_refund_from, trial_refund_to = self.get_trial_days(
+            trial_info = self.get_trial_days(
                 trial_start_date,
                 trial_end_date,
             )
-            trial_amount = sum(daily_expenses.get(day, Decimal("0")) for day in trial_days)
-
+            trial_amount = sum(
+                daily_expenses.get(day, Decimal("0")) for day in trial_info.trial_days
+            )
+            trial_days = trial_info.trial_days
             refund_lines.append(
                 Refund(
                     trial_amount,
-                    trial_refund_from,
-                    trial_refund_to,
+                    trial_info.refund_from,
+                    trial_info.refund_to,
                     (
                         "Refund due to trial period "
                         f"(from {trial_start_date.strftime("%d %b %Y")} "  # type: ignore
@@ -614,7 +618,7 @@ async def generate_refund_lines(
         linked_datasource_type: str,
         datasource_id: str,
         datasource_name: str,
-    ):
+    ) -> list[str]:
         """
         This function calculates the refund lines for a billing period
         """
@@ -680,24 +684,39 @@ def get_trial_days(
         self,
         trial_start_date: date | None,
         trial_end_date: date | None,
-    ):
+    ) -> TrialInfo:
         # Trial period can start or end on month other than billing month period.
         # In this situation, we need to limit refunded expenses
         # to a period overlapping with billing month.
         # Example, billing month is June 1-30, a trial period is May 17 - June 17
         # We need to refund expenses from June 1st to June 17th
         if not trial_start_date and not trial_end_date:
-            return None, None, None
+            return TrialInfo()
         trial_refund_from = max(trial_start_date, self.billing_start_date.date())
         trial_refund_to = min(trial_end_date, self.billing_end_date.date())
         trial_days = {
             dt.date().day for dt in rrule(DAILY, dtstart=trial_refund_from, until=trial_refund_to)
         }
-        return trial_days, trial_refund_from, trial_refund_to
+        return TrialInfo(trial_days, trial_refund_from, trial_refund_to)
 
     def get_entitlement_days(
-        self, entitlement_start_date: str, entitlement_end_date: str, trial_days: set[int]
-    ):
+        self,
+        entitlement_start_date: str,
+        entitlement_end_date: str | None = None,
+        trial_days: set[int] | None = None,
+    ) -> set[int]:
+        """
+        Calculate the set of entitlement day numbers within a date range,
+        excluding the given trial day numbers.
+
+        Args:
+            entitlement_start_date (str): ISO format start date string.
+            entitlement_end_date (Optional[str]): ISO format end date string, or None.
+            trial_days (Optional[Set[int]]): Set of trial day numbers (1–31)
+
+        Returns:
+            Set[int]: Set of day numbers (1–31) representing entitlement days.
+        """
         start_date = max(datetime.fromisoformat(entitlement_start_date), self.billing_start_date)
         end_date = min(
             datetime.fromisoformat(entitlement_end_date)
@@ -722,7 +741,7 @@ def generate_charge_line(
         price: Decimal,
         datasource_name: str,
         description: str = "",
-    ):
+    ) -> str:
         """
         This function generates a charge line for a vendor and datasource.
         """

ffc/management/commands/process_billing.py

@@ -4,7 +4,7 @@
 from dateutil.relativedelta import relativedelta
 from django.core.management.base import BaseCommand
 
-from ffc.process_billing import (
+from ffc.billing.process_billing import (
     process_billing,
 )
 

tests/conftest.py

@@ -1,14 +1,17 @@
 import copy
 from datetime import UTC, datetime
 from decimal import Decimal
+from typing import Generator
 
+import httpx
 import jwt
 import pytest
 import responses
 from swo.mpt.extensions.runtime.djapp.conf import get_for_product
 
 from ffc.billing.dataclasses import AuthorizationProcessResult
-from ffc.process_billing import AuthorizationProcessor
+from ffc.billing.process_billing import AuthorizationProcessor
+from ffc.clients.base import BaseAsyncAPIClient
 
 
 @pytest.fixture()
@@ -2665,3 +2668,30 @@ def patch_fetch_organizations_agr_000(
         "fetch_organizations",
         return_value=org_mock_generator_agr_000,
     )
+
+
+class TestClientAuth(httpx.Auth):
+    def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
+        request.headers["Authorization"] = "Bearer fake token"
+        yield request
+
+
+class FakeAPIClient(BaseAsyncAPIClient):
+    @property
+    def base_url(self) -> str:
+        return "https://local.local/v1"
+
+    @property
+    def auth(self):
+        return TestClientAuth()
+
+    def get_pagination_meta(self, response):
+        return response["meta"]["pagination"]
+
+    def get_page_data(self, response):
+        return response["data"]
+
+
+@pytest.fixture()
+def fake_apiclient():
+    return FakeAPIClient(limit=2)

tests/ffc/test_base_client.py

@@ -1,37 +1,11 @@
-from collections.abc import Generator
-
 import httpx
 import pytest
 
 from ffc.clients.base import BaseAsyncAPIClient, PaginationSupportMixin
 
 
-class TestClientAuth(httpx.Auth):
-    def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
-        request.headers["Authorization"] = "Bearer fake token"
-        yield request
-
-
-class TestClient(BaseAsyncAPIClient):
-    @property
-    def base_url(self) -> str:
-        return "https://local.local/v1"
-
-    @property
-    def auth(self):
-        return TestClientAuth()
-
-    def get_pagination_meta(self, response):
-        return response["meta"]["pagination"]
-
-    def get_page_data(self, response):
-        return response["data"]
-
-
 @pytest.mark.asyncio()
-async def test_collection_iterator_paginates(httpx_mock):
-    client = TestClient(limit=2)
-
+async def test_collection_iterator_paginates(httpx_mock, fake_apiclient):
     endpoint = "/catalog/authorizations"
     rql = "eq(mytestfield,'value')"
 
@@ -60,7 +34,7 @@ async def test_collection_iterator_paginates(httpx_mock):
     )
 
     items = []
-    async for item in client.collection_iterator(endpoint, rql=rql):
+    async for item in fake_apiclient.collection_iterator(endpoint, rql=rql):
         items.append(item)
 
     assert [item["id"] for item in items] == ["AUT-1111-1111", "AUT-2222-2222", "AUT-3333-3333"]
@@ -74,14 +48,12 @@ async def test_collection_iterator_paginates(httpx_mock):
     for r in reqs:
         assert r.headers["Authorization"] == "Bearer fake token"
 
-    await client.close()
-    assert client.httpx_client.is_closed
+    await fake_apiclient.close()
+    assert fake_apiclient.httpx_client.is_closed
 
 
 @pytest.mark.asyncio()
-async def test_collection_iterator_paginates_404(httpx_mock):
-    client = TestClient(limit=2)
-
+async def test_collection_iterator_paginates_404(httpx_mock, fake_apiclient):
     endpoint = "/catalog/authorizations"
     rql = "eq(mytestfield,'value')"
 
@@ -92,14 +64,14 @@ async def test_collection_iterator_paginates_404(httpx_mock):
     )
 
     with pytest.raises(httpx.HTTPStatusError):
-        await anext(client.collection_iterator(endpoint, rql=rql))
+        await anext(fake_apiclient.collection_iterator(endpoint, rql=rql))
     [req] = httpx_mock.get_requests()
     assert req.method == "GET"
     assert str(req.url) == f"https://local.local/v1/catalog/authorizations?{rql}&limit=2&offset=0"
     assert req.headers["Authorization"] == "Bearer fake token"
 
-    await client.close()
-    assert client.httpx_client.is_closed
+    await fake_apiclient.close()
+    assert fake_apiclient.httpx_client.is_closed
 
 
 def test_cannot_instantiate_paginationsupportmixin_directly():