File diff c5dd7984bcc1 → 786291902231
tests/testutil.py
Show inline comments
 
new file 100644
 
"""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
 

	
 
FY_START_DATE = datetime.date(2020, 3, 1)
 
FY_MID_DATE = datetime.date(2020, 9, 1)
 

	
 
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):
 
    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)