conservancy_beancount/tests/testutil.py

"""Mock Beancount objects for testing"""
# 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/>.

import datetime

import beancount.core.amount as bc_amount
import beancount.core.data as bc_data

from decimal import Decimal

EXTREME_FUTURE_DATE = datetime.date(datetime.MAXYEAR, 12, 30)
FUTURE_DATE = datetime.date.today() + datetime.timedelta(days=365 * 99)
FY_START_DATE = datetime.date(2020, 3, 1)
FY_MID_DATE = datetime.date(2020, 9, 1)
PAST_DATE = datetime.date(2000, 1, 1)

def check_post_meta(txn, *expected_meta, default=None):
    assert len(txn.postings) == len(expected_meta)
    for post, expected in zip(txn.postings, expected_meta):
        if not expected:
            assert not post.meta
        else:
            assert all(post.meta.get(key, default) == value
                       for key, value in expected.items())

def parse_date(s, fmt='%Y-%m-%d'):
    return datetime.datetime.strptime(s, fmt).date()

def Posting(account, number,
            currency='USD', cost=None, price=None, flag=None,
            **meta):
    if not meta:
        meta = None
    return bc_data.Posting(
        account,
        bc_amount.Amount(Decimal(number), currency),
        cost,
        price,
        flag,
        meta,
    )

class Transaction:
    def __init__(self,
                 date=FY_MID_DATE, flag='*', payee=None,
                 narration='', tags=None, links=None, postings=None,
                 **meta):
        if isinstance(date, str):
            date = parse_date(date)
        self.date = date
        self.flag = flag
        self.payee = payee
        self.narration = narration
        self.tags = set(tags or '')
        self.links = set(links or '')
        self.postings = []
        self.meta = {
            'filename': '<test>',
            'lineno': 0,
        }
        self.meta.update(meta)
        for posting in postings:
            self.add_posting(*posting)

    def add_posting(self, arg, *args, **kwargs):
        """Add a posting to this transaction. Use any of these forms:

           txn.add_posting(account, number, …, kwarg=value, …)
           txn.add_posting(account, number, …, posting_kwargs_dict)
           txn.add_posting(posting_object)
        """
        if kwargs:
            posting = Posting(arg, *args, **kwargs)
        elif args:
            if isinstance(args[-1], dict):
                kwargs = args[-1]
                args = args[:-1]
            posting = Posting(arg, *args, **kwargs)
        else:
            posting = arg
        self.postings.append(posting)