"""Test data.balance_of function"""
# Copyright © 2020  Brett Smith
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from decimal import Decimal
from operator import methodcaller

import pytest

from . import testutil

from conservancy_beancount import data

is_cash_eq = data.Account.is_cash_equivalent

@pytest.fixture
def payable_payment_txn():
    return testutil.Transaction(postings=[
        ('Liabilities:Payable:Accounts', 50),
        ('Assets:Checking', -50),
        ('Expenses:BankingFees', 5),
        ('Assets:Checking', -5),
    ])

def balance_under(txn, *accts):
    pred = methodcaller('is_under', *accts)
    return data.balance_of(txn, pred)

def test_balance_of_simple_txn():
    txn = testutil.Transaction(postings=[
        ('Assets:Cash', 50),
        ('Income:Donations', -50),
    ])
    assert balance_under(txn, 'Assets') == 50
    assert balance_under(txn, 'Income') == -50

def test_zero_balance_of(payable_payment_txn):
    assert balance_under(payable_payment_txn, 'Equity') == 0
    assert balance_under(payable_payment_txn, 'Assets:Cash') == 0
    assert balance_under(payable_payment_txn, 'Liabilities:CreditCard') == 0

def test_nonzero_balance_of(payable_payment_txn):
    assert balance_under(payable_payment_txn, 'Assets', 'Expenses') == -50
    assert balance_under(payable_payment_txn, 'Assets', 'Liabilities') == -5

def test_multiarg_balance_of():
    txn = testutil.Transaction(postings=[
        ('Liabilities:CreditCard', 650),
        ('Expenses:BankingFees', 5),
        ('Assets:Checking', -655),
    ])
    assert data.balance_of(txn, is_cash_eq, data.Account.is_credit_card) == -5

def test_balance_of_multipost_txn(payable_payment_txn):
    assert data.balance_of(payable_payment_txn, is_cash_eq) == -55