"""Test PaymentBatchProcessor.""" from decimal import Decimal from typing import List from unittest.mock import Mock, patch from src.constants import TaxCorrectionTypes from src.models import FlowthroughAllocation, TaxCorrection, TaxCorrectionVAT from src.processors.base import payment_batch_processor as processor from src.processors.base.worksheet_calculator import WorksheetCalculator from src.processors.models import ContractLevelFlowthroughAllocationData from tests.unit.factories import ( AccountFactory, ContractCloseBalanceFactory, EligibleAccountLevelDataFactory, EventFactory, FlowthroughAllocationFactory, PaginatedPayableBalanceAfterTaxEntriesFactory, PayableBalanceAfterTaxEntryFactory, TaxCorrectionFactory, TaxCorrectionVATFactory, ) mock_eligible_account = EligibleAccountLevelDataFactory.build( account_id=1, currency_code='USD', country_of_tax_policy='USA' ) mock_event = EventFactory.build(statment_period_id=318) mock_close_balance = ContractCloseBalanceFactory.build( account_id=mock_eligible_account.account_id, amount=Decimal('1000.50') ) mock_wht_correction = TaxCorrectionFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance.contract_id, amount=Decimal('-100.10'), ) mock_vat_correction = TaxCorrectionVATFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance.contract_id, vat_amount_payee_currency=Decimal('50.00'), ) mock_flowthrough_allocation = FlowthroughAllocationFactory.build( contract_id=mock_close_balance.contract_id, amount_to_payment=Decimal('200.00'), ) class TestPaymentBatchProcessor: """PaymentBatchProcessor test suite.""" def test__init__(self) -> None: """test init.""" test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) assert test_processor._abacus_event == mock_event assert test_processor._statement_period_id == mock_event.statement_period_id assert test_processor._accounts_by_id == {1: mock_eligible_account} assert test_processor._contract_worksheets_by_id == {} @patch.object(processor, 'fetch_all_contract_closing_balance_entries') def test__get_closing_balance_related_data( self, mock_fetch_all_contract_closing_balance_entries: Mock ) -> None: """Test retrieve closing balance success.""" mock_fetch_all_contract_closing_balance_entries.return_value = [ mock_close_balance ] test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) assert test_processor._get_contract_ids() == [] response = test_processor._get_closing_balance_related_data() mock_fetch_all_contract_closing_balance_entries.assert_called_once_with( mock_event.statement_period_id, [ AccountFactory.build( account_id=mock_eligible_account.account_id, country_of_tax_residence=None, currency_code=None, payment_entity_id=None, ) ], ) assert response == [mock_close_balance] @patch.object(processor.PaymentBatchProcessor, '_get_closing_balance_related_data') def test__setup_contract_worksheets( self, mock_get_closing_balance_related_data: Mock, ) -> None: """Test worksheet setup.""" mock_get_closing_balance_related_data.return_value = [mock_close_balance] test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._setup_contract_worksheets() mock_get_closing_balance_related_data.assert_called_once() response_worksheet = test_processor._contract_worksheets_by_id[ mock_close_balance.contract_id ] assert type(response_worksheet).__name__ == 'WorksheetCalculator' assert response_worksheet.account_id == mock_eligible_account.account_id assert response_worksheet.contract_id == mock_close_balance.contract_id assert response_worksheet.payable_amount_pre_tax == mock_close_balance.amount assert response_worksheet.tax_withholding_amount is None assert response_worksheet.vat_amount is None assert response_worksheet.payable_amount_post_tax == mock_close_balance.amount def test__setup_contract_worksheets_builds_statement_period_map( self, ) -> None: """Test worksheet setup always builds the contract statement period mapping.""" with patch.object( processor.PaymentBatchProcessor, '_get_closing_balance_related_data' ) as mock_get_closing_balance_related_data: mock_get_closing_balance_related_data.return_value = [mock_close_balance] test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._setup_contract_worksheets() assert test_processor._contract_statement_periods == { mock_close_balance.contract_id: mock_close_balance.statement_period_id } @patch.object(processor, 'get_contract_balance_after_tax_entries') def test__get_created_worksheets( self, mock_get_contract_balance_after_tax_entries: Mock ) -> None: """Test get newly created worksheets.""" mock_worksheet = WorksheetCalculator(mock_close_balance, mock_eligible_account) mock_balance_after_tax = PayableBalanceAfterTaxEntryFactory.build( contract_id=mock_close_balance.contract_id ) mock_get_contract_balance_after_tax_entries.return_value = ( PaginatedPayableBalanceAfterTaxEntriesFactory.build( items=[mock_balance_after_tax] ) ) test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._contract_worksheets_by_id[mock_close_balance.contract_id] = ( mock_worksheet ) response = test_processor._get_created_worksheets() assert response == [mock_balance_after_tax] mock_get_contract_balance_after_tax_entries.assert_called_once_with( mock_event.abacus_event_id, limit=100, contract_ids=[mock_close_balance.contract_id], ) @patch.object(processor, 'bulk_create_contract_payable_details') @patch.object(processor, 'get_contract_balance_after_tax_entries') @patch.object(processor, 'bulk_create_worksheet_contract_balance_after_tax') def test__post_contract_worksheets_and_details( self, mock_bulk_create_worksheet_contract_balance_after_tax: Mock, mock_get_contract_balance_after_tax_entries: Mock, mock_bulk_create_contract_payable_details: Mock, ) -> None: """Test bulk creation of worksheets and details.""" mock_worksheet = WorksheetCalculator(mock_close_balance, mock_eligible_account) mock_balance_after_tax = PayableBalanceAfterTaxEntryFactory.build( contract_id=mock_close_balance.contract_id ) mock_bulk_create_worksheet_contract_balance_after_tax.return_value = None mock_get_contract_balance_after_tax_entries.return_value = ( PaginatedPayableBalanceAfterTaxEntriesFactory.build( items=[mock_balance_after_tax] ) ) test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._contract_worksheets_by_id[mock_close_balance.contract_id] = ( mock_worksheet ) test_processor._post_contract_worksheets_and_details() mock_bulk_create_worksheet_contract_balance_after_tax.assert_called_once_with( mock_event.abacus_event_id, mock_event.statement_period_id, [mock_worksheet.get_payable_balance_after_tax_item()], ) mock_get_contract_balance_after_tax_entries.assert_called_once_with( mock_event.abacus_event_id, limit=100, contract_ids=[mock_close_balance.contract_id], ) mock_bulk_create_contract_payable_details.assert_called_once_with( mock_event.abacus_event_id, mock_event.statement_period_id, mock_worksheet.get_payable_detail_items( mock_balance_after_tax.worksheet_account_contract_payable_after_tax_id ), ) @patch.object(processor, 'fetch_all_pending_tax_corrections') def test__append_wht_corrections_to_worksheets_with_tap_flowthrough( self, mock_fetch_all_pending_tax_corrections: Mock, ) -> None: """Test fetch and append WHT corrections grouped per statement period.""" # Create two closing balances with different statement periods mock_close_balance_current = ContractCloseBalanceFactory.build( account_id=mock_eligible_account.account_id, amount=Decimal('1000.00'), contract_id=1, statement_period_id=318, ) mock_close_balance_previous = ContractCloseBalanceFactory.build( account_id=mock_eligible_account.account_id, amount=Decimal('2000.00'), contract_id=2, statement_period_id=317, ) # Create corrections for both contracts mock_wht_correction_current = TaxCorrectionFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance_current.contract_id, amount=Decimal('-100.00'), ) mock_wht_correction_previous = TaxCorrectionFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance_previous.contract_id, amount=Decimal('-200.00'), ) # Mock to return different corrections for different statement periods def fetch_corrections_side_effect( correction_type: TaxCorrectionTypes, contract_ids: List[int], statement_period_id: int, ) -> List[TaxCorrection]: if statement_period_id == 318: return [mock_wht_correction_current] elif statement_period_id == 317: return [mock_wht_correction_previous] return [] mock_fetch_all_pending_tax_corrections.side_effect = ( fetch_corrections_side_effect ) test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._contract_worksheets_by_id[ mock_close_balance_current.contract_id ] = WorksheetCalculator(mock_close_balance_current, mock_eligible_account) test_processor._contract_worksheets_by_id[ mock_close_balance_previous.contract_id ] = WorksheetCalculator(mock_close_balance_previous, mock_eligible_account) test_processor._contract_statement_periods = { mock_close_balance_current.contract_id: mock_close_balance_current.statement_period_id, mock_close_balance_previous.contract_id: mock_close_balance_previous.statement_period_id, } test_processor._append_wht_corrections_to_worksheets() # Verify that corrections were fetched for both statement periods assert mock_fetch_all_pending_tax_corrections.call_count == 2 # Verify corrections were applied to correct worksheets worksheet_current = test_processor._contract_worksheets_by_id[ mock_close_balance_current.contract_id ] assert ( worksheet_current.tax_withholding_amount == mock_wht_correction_current.amount ) worksheet_previous = test_processor._contract_worksheets_by_id[ mock_close_balance_previous.contract_id ] assert ( worksheet_previous.tax_withholding_amount == mock_wht_correction_previous.amount ) @patch.object(processor, 'fetch_all_pending_tax_corrections') def test__append_wht_corrections_to_worksheets_multiple_periods( self, mock_fetch_all_pending_tax_corrections: Mock, ) -> None: """Test fetch and append WHT corrections across multiple statement periods.""" # Create two closing balances with different statement periods mock_close_balance_current = ContractCloseBalanceFactory.build( account_id=mock_eligible_account.account_id, amount=Decimal('1000.00'), contract_id=1, statement_period_id=318, ) mock_close_balance_previous = ContractCloseBalanceFactory.build( account_id=mock_eligible_account.account_id, amount=Decimal('2000.00'), contract_id=2, statement_period_id=317, ) # Create corrections for both contracts mock_wht_correction_current = TaxCorrectionFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance_current.contract_id, amount=Decimal('-100.00'), ) mock_wht_correction_previous = TaxCorrectionFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance_previous.contract_id, amount=Decimal('-200.00'), ) # Mock to return different corrections for different statement periods def fetch_corrections_side_effect( correction_type: TaxCorrectionTypes, contract_ids: List[int], statement_period_id: int, ) -> List[TaxCorrection]: if statement_period_id == 318: return [mock_wht_correction_current] elif statement_period_id == 317: return [mock_wht_correction_previous] return [] mock_fetch_all_pending_tax_corrections.side_effect = ( fetch_corrections_side_effect ) test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._contract_worksheets_by_id[ mock_close_balance_current.contract_id ] = WorksheetCalculator(mock_close_balance_current, mock_eligible_account) test_processor._contract_worksheets_by_id[ mock_close_balance_previous.contract_id ] = WorksheetCalculator(mock_close_balance_previous, mock_eligible_account) test_processor._contract_statement_periods = { mock_close_balance_current.contract_id: mock_close_balance_current.statement_period_id, mock_close_balance_previous.contract_id: mock_close_balance_previous.statement_period_id, } test_processor._append_wht_corrections_to_worksheets() # Verify that corrections were fetched for both statement periods assert mock_fetch_all_pending_tax_corrections.call_count == 2 # Verify corrections were applied to correct worksheets worksheet_current = test_processor._contract_worksheets_by_id[ mock_close_balance_current.contract_id ] assert ( worksheet_current.tax_withholding_amount == mock_wht_correction_current.amount ) worksheet_previous = test_processor._contract_worksheets_by_id[ mock_close_balance_previous.contract_id ] assert ( worksheet_previous.tax_withholding_amount == mock_wht_correction_previous.amount ) @patch.object(processor, 'bulk_create_worksheet_contract_balance_after_tax') def test__post_contract_worksheets_and_details_none( self, mock_bulk_create_worksheet_contract_balance_after_tax: Mock ) -> None: """Test bulk creation skipped when no worksheets exist.""" test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._contract_worksheets_by_id = {} test_processor._post_contract_worksheets_and_details() mock_bulk_create_worksheet_contract_balance_after_tax.assert_not_called() @patch.object(processor, 'fetch_all_pending_tax_corrections_vat') def test__append_vat_corrections_to_worksheets( self, mock_fetch_all_pending_tax_corrections_vat: Mock, ) -> None: """Test fetch and append VAT corrections grouped per statement period.""" # Create two closing balances with different statement periods mock_close_balance_current = ContractCloseBalanceFactory.build( account_id=mock_eligible_account.account_id, amount=Decimal('1000.00'), contract_id=1, statement_period_id=318, ) mock_close_balance_previous = ContractCloseBalanceFactory.build( account_id=mock_eligible_account.account_id, amount=Decimal('2000.00'), contract_id=2, statement_period_id=317, ) # Create VAT corrections for both contracts mock_vat_correction_current = TaxCorrectionVATFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance_current.contract_id, vat_amount_payee_currency=Decimal('100.00'), ) mock_vat_correction_previous = TaxCorrectionVATFactory.build( account_id=mock_eligible_account.account_id, contract_id=mock_close_balance_previous.contract_id, vat_amount_payee_currency=Decimal('200.00'), ) # Mock to return different corrections for different statement periods def fetch_corrections_side_effect( contract_ids: List[int], statement_period_id: int, ) -> List[TaxCorrectionVAT]: if statement_period_id == 318: return [mock_vat_correction_current] elif statement_period_id == 317: return [mock_vat_correction_previous] return [] mock_fetch_all_pending_tax_corrections_vat.side_effect = ( fetch_corrections_side_effect ) test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._contract_worksheets_by_id[ mock_close_balance_current.contract_id ] = WorksheetCalculator(mock_close_balance_current, mock_eligible_account) test_processor._contract_worksheets_by_id[ mock_close_balance_previous.contract_id ] = WorksheetCalculator(mock_close_balance_previous, mock_eligible_account) test_processor._contract_statement_periods = { mock_close_balance_current.contract_id: mock_close_balance_current.statement_period_id, mock_close_balance_previous.contract_id: mock_close_balance_previous.statement_period_id, } test_processor._append_vat_corrections_to_worksheets() # Verify that corrections were fetched for both statement periods assert mock_fetch_all_pending_tax_corrections_vat.call_count == 2 # Verify corrections were applied to correct worksheets worksheet_current = test_processor._contract_worksheets_by_id[ mock_close_balance_current.contract_id ] assert worksheet_current.vat_amount == mock_vat_correction_current.amount worksheet_previous = test_processor._contract_worksheets_by_id[ mock_close_balance_previous.contract_id ] assert worksheet_previous.vat_amount == mock_vat_correction_previous.amount @patch.object(processor, 'fetch_all_pending_tax_corrections_vat') def test__append_vat_corrections_disabled_via_parameter_and_flag( self, mock_fetch_all_pending_tax_corrections_vat: Mock, ) -> None: """Test VAT corrections are not processed when append_vat_corrections=False (check payments).""" test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account], append_vat_corrections=False ) test_processor._contract_worksheets_by_id[mock_close_balance.contract_id] = ( WorksheetCalculator(mock_close_balance, mock_eligible_account) ) test_processor._contract_statement_periods = { mock_close_balance.contract_id: mock_event.statement_period_id } pre_response_worksheet = test_processor._contract_worksheets_by_id[ mock_close_balance.contract_id ] assert pre_response_worksheet.vat_amount is None test_processor._append_vat_corrections_to_worksheets() # VAT corrections should not be fetched when append_vat_corrections=False mock_fetch_all_pending_tax_corrections_vat.assert_not_called() post_response_worksheet = test_processor._contract_worksheets_by_id[ mock_close_balance.contract_id ] # VAT amount should remain None assert post_response_worksheet.vat_amount is None @patch.object(processor, 'fetch_all_flowthrough_allocation_entries') def test__get_flowthrough_allocation_by_contract( self, mock_fetch_all_flowthrough_allocation_entries: Mock ) -> None: """Test get flowthrough allocation data for contracts.""" allocation_1 = FlowthroughAllocationFactory.build( contract_id=mock_close_balance.contract_id, amount_to_payment=Decimal('100.00'), ) allocation_2 = FlowthroughAllocationFactory.build( contract_id=mock_close_balance.contract_id, amount_to_payment=Decimal('150.00'), ) mock_fetch_all_flowthrough_allocation_entries.return_value = [ allocation_1, allocation_2, ] test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) test_processor._contract_worksheets_by_id[mock_close_balance.contract_id] = ( WorksheetCalculator(mock_close_balance, mock_eligible_account) ) result = test_processor._get_flowthrough_allocation_by_contract() mock_fetch_all_flowthrough_allocation_entries.assert_called_once_with( [mock_close_balance.contract_id] ) assert mock_close_balance.contract_id in result contract_data = result[mock_close_balance.contract_id] assert isinstance(contract_data, ContractLevelFlowthroughAllocationData) assert contract_data.sum == Decimal('250.00') assert contract_data.items == [allocation_1, allocation_2] @patch.object( processor.PaymentBatchProcessor, '_get_flowthrough_allocation_by_contract' ) def test__append_flowthrough_allocations_to_worksheets( self, mock_get_flowthrough_allocation: Mock ) -> None: """Test append flowthrough allocations to worksheets.""" contract_data = ContractLevelFlowthroughAllocationData( sum=Decimal('200.00'), items=[mock_flowthrough_allocation] ) mock_get_flowthrough_allocation.return_value = { mock_close_balance.contract_id: contract_data } test_processor = processor.PaymentBatchProcessor( mock_event, [mock_eligible_account] ) mock_worksheet = Mock(spec=WorksheetCalculator) test_processor._contract_worksheets_by_id[mock_close_balance.contract_id] = ( mock_worksheet ) test_processor._append_flowthrough_allocations_to_worksheets() mock_get_flowthrough_allocation.assert_called_once() mock_worksheet.apply_flowthrough_items.assert_called_once_with(contract_data)