conservancy_beancount/tests/test_reports_rewrite.py

"""test_reports_rewrite - Unit tests for report rewrite functionality"""
# Copyright © 2020  Brett Smith
# License: AGPLv3-or-later WITH Beancount-Plugin-Additional-Permission-1.0
#
# Full copyright and licensing details can be found at toplevel file
# LICENSE.txt in the repository.

import datetime

import pytest

from decimal import Decimal

import yaml

from . import testutil

from conservancy_beancount import data
from conservancy_beancount.reports import rewrite

CMP_OPS = frozenset('< <= == != >= >'.split())

@pytest.mark.parametrize('name', ['Equity:Other', 'Expenses:Other', 'Income:Other'])
@pytest.mark.parametrize('operator', CMP_OPS)
def test_account_condition(name, operator):
    operand = 'Expenses:Other'
    txn = testutil.Transaction(postings=[(name, -5)])
    post, = data.Posting.from_txn(txn)
    tester = rewrite.AccountTest(operator, operand)
    assert tester(post) == eval(f'name {operator} operand')

@pytest.mark.parametrize('name,expected', [
    ('Expenses:Postage', True),
    ('Expenses:Tax', True),
    ('Expenses:Tax:Sales', True),
    ('Expenses:Tax:VAT', True),
    ('Expenses:Taxes', False),
    ('Expenses:Other', False),
    ('Liabilities:Tax', False),
])
def test_account_in_condition(name, expected):
    txn = testutil.Transaction(postings=[(name, 5)])
    post, = data.Posting.from_txn(txn)
    tester = rewrite.AccountTest('in', 'Expenses:Tax Expenses:Postage')
    assert tester(post) == expected

@pytest.mark.parametrize('n', range(3, 12, 3))
@pytest.mark.parametrize('operator', CMP_OPS)
def test_date_condition(n, operator):
    date = datetime.date(2020, n, n)
    txn = testutil.Transaction(date=date, postings=[
        ('Income:Other', -5),
    ])
    post, = data.Posting.from_txn(txn)
    tester = rewrite.DateTest(operator, '2020-06-06')
    assert tester(post) == eval(f'n {operator} 6')

@pytest.mark.parametrize('value', ['test', 'testvalue', 'testzed'])
@pytest.mark.parametrize('operator', CMP_OPS)
def test_metadata_condition(value, operator):
    key = 'testkey'
    operand = 'testvalue'
    txn = testutil.Transaction(postings=[
        ('Income:Other', -5, {key: value}),
    ])
    post, = data.Posting.from_txn(txn)
    tester = rewrite.MetadataTest(key, operator, operand)
    assert tester(post) == eval(f'value {operator} operand')

@pytest.mark.parametrize('value,expected', [
    ('Root', False),
    ('Root:Branch', True),
    ('Root:Branch:Leaf', True),
    ('Branch', False),
    ('RootBranch:Leaf', False),
    (None, False),
])
def test_metadata_in_condition(value, expected):
    key = 'testkey'
    meta = {} if value is None else {key: value}
    txn = testutil.Transaction(postings=[
        ('Income:Other', -5, meta),
    ])
    post, = data.Posting.from_txn(txn)
    tester = rewrite.MetadataTest(key, 'in', 'Root:Branch')
    assert tester(post) == expected

@pytest.mark.parametrize('value', ['4.5', '4.75', '5'])
@pytest.mark.parametrize('operator', CMP_OPS)
def test_number_condition(value, operator):
    operand = '4.75'
    txn = testutil.Transaction(postings=[
        ('Expenses:Other', value),
    ])
    post, = data.Posting.from_txn(txn)
    tester = rewrite.NumberTest(operator, operand)
    assert tester(post) == eval(f'value {operator} operand')

@pytest.mark.parametrize('subject,operand', [
    ('.account', 'Income:Other'),
    ('.date', '1990-05-10'),
    ('.number', '5.79'),
    ('testkey', 'testvalue'),
])
@pytest.mark.parametrize('operator', CMP_OPS)
def test_parse_good_condition(subject, operator, operand):
    actual = rewrite.TestRegistry.parse(f'{subject}{operator}{operand}')
    if subject == '.account':
        assert isinstance(actual, rewrite.AccountTest)
        assert actual.operand == operand
    elif subject == '.date':
        assert isinstance(actual, rewrite.DateTest)
        assert actual.operand == datetime.date(1990, 5, 10)
    elif subject == '.number':
        assert isinstance(actual, rewrite.NumberTest)
        assert actual.operand == Decimal(operand)
    else:
        assert isinstance(actual, rewrite.MetadataTest)
        assert actual.key == 'testkey'
        assert actual.operand == 'testvalue'

@pytest.mark.parametrize('cond_s', [
    '.account = Income:Other',  # Bad operator
    '.account===Equity:Other',  # Bad operand (`=Equity:Other` is not valid)
    '.account in foo',  # Bad operand
    '.date == 1990-90-5',  # Bad operand
    '.date in 1990-9-9',  # Bad operator
    '.number > 0xff',  # Bad operand
    '.number in 16',  # Bad operator
    'units.number == 5',  # Bad subject (syntax)
    '.units == 5',  # Bad subject (unknown)
])
def test_parse_bad_condition(cond_s):
    with pytest.raises(ValueError):
        rewrite.TestRegistry.parse(cond_s)

@pytest.mark.parametrize('value', ['Equity:Other', 'Income:Other'])
def test_account_set(value):
    value = data.Account(value)
    txn = testutil.Transaction(postings=[
        ('Expenses:Other', 5),
    ])
    post, = data.Posting.from_txn(txn)
    setter = rewrite.AccountSet('=', value)
    assert setter(post) == ('account', value)

@pytest.mark.parametrize('key', ['aa', 'bb'])
def test_metadata_set(key):
    txn_meta = {'filename': 'metadata_set', 'lineno': 100}
    post_meta = {'aa': 'one', 'bb': 'two'}
    meta = {'aa': 'one', 'bb': 'two'}
    txn = testutil.Transaction(**txn_meta, postings=[
        ('Expenses:Other', 5, post_meta),
    ])
    post, = data.Posting.from_txn(txn)
    setter = rewrite.MetadataSet(key, '=', 'newvalue')
    assert setter(post) == (key, 'newvalue')

@pytest.mark.parametrize('value', ['0.25', '-.5', '1.9'])
@pytest.mark.parametrize('currency', ['USD', 'EUR', 'INR'])
def test_number_set(value, currency):
    txn = testutil.Transaction(postings=[
        ('Expenses:Other', 5, currency),
    ])
    post, = data.Posting.from_txn(txn)
    setter = rewrite.NumberSet('*=', value)
    assert setter(post) == ('units', testutil.Amount(Decimal(value) * 5, currency))

@pytest.mark.parametrize('subject,operator,operand', [
    ('.account', '=', 'Income:Other'),
    ('.number', '*=', '.5'),
    ('.number', '*=', '-1'),
    ('testkey', '=', 'testvalue'),
])
def test_parse_good_set(subject, operator, operand):
    actual = rewrite.SetRegistry.parse(f'{subject}{operator}{operand}')
    if subject == '.account':
        assert isinstance(actual, rewrite.AccountSet)
        assert actual.value == operand
    elif subject == '.number':
        assert isinstance(actual, rewrite.NumberSet)
        assert actual.value == Decimal(operand)
    else:
        assert isinstance(actual, rewrite.MetadataSet)
        assert actual.key == subject
        assert actual.value == operand

@pytest.mark.parametrize('set_s', [
    '.account==Equity:Other',  # Bad operand (`=Equity:Other` is not valid)
    '.account*=2',  # Bad operator
    '.date = 2020-02-20',  # Bad subject
    '.number*=0xff',  # Bad operand
    '.number=5',  # Bad operator
    'testkey += foo',  # Bad operator
    'testkey *= 3',  # Bad operator
])
def test_parse_bad_set(set_s):
    with pytest.raises(ValueError):
        rewrite.SetRegistry.parse(set_s)

def test_good_rewrite_rule():
    rule = rewrite.RewriteRule({
        'if': ['.account in Income'],
        'add': ['income-type = Other'],
    })
    txn = testutil.Transaction(postings=[
        ('Assets:Cash', 10),
        ('Income:Other', -10),
    ])
    cash_post, inc_post = data.Posting.from_txn(txn)
    assert not rule.match(cash_post)
    assert rule.match(inc_post)
    new_post, = rule.rewrite(inc_post)
    assert new_post.account == 'Income:Other'
    assert new_post.units == testutil.Amount(-10)
    assert new_post.meta.pop('income-type', None) == 'Other'
    assert new_post.meta
    assert new_post.meta.date == txn.date

def test_complicated_rewrite_rule():
    account = 'Income:Donations'
    income_key = 'income-type'
    income_type = 'Refund'
    rule = rewrite.RewriteRule({
        'if': ['.account == Expenses:Refunds'],
        'project': [
            f'.account = {account}',
            '.number *= .8',
            f'{income_key} = {income_type}',
        ],
        'general': [
            f'.account = {account}',
            '.number *= .2',
            f'{income_key} = {income_type}',
            'project = Conservancy',
        ],
    })
    txn = testutil.Transaction(postings=[
        ('Assets:Cash', -12),
        ('Expenses:Refunds', 12, {'project': 'Bravo'}),
    ])
    cash_post, src_post = data.Posting.from_txn(txn)
    assert not rule.match(cash_post)
    assert rule.match(src_post)
    proj_post, gen_post = rule.rewrite(src_post)
    assert proj_post.account == 'Income:Donations'
    assert proj_post.units == testutil.Amount('9.60')
    assert proj_post.meta[income_key] == income_type
    assert proj_post.meta['project'] == 'Bravo'
    assert gen_post.account == 'Income:Donations'
    assert gen_post.units == testutil.Amount('2.40')
    assert gen_post.meta[income_key] == income_type
    assert gen_post.meta['project'] == 'Conservancy'

@pytest.mark.parametrize('source', [
    # Account assignments
    {'if': ['.account in Income Expenses'], 'then': ['.account = Equity']},
    {'if': ['.account == Assets:PettyCash'], 'then': ['.account = Assets:Cash']},
    {'if': ['.account == Liabilities:CreditCard'], 'then': ['.account = Liabilities:Visa']},
    # Number splits
    {'if': ['.date >= 2020-01-01'], 'a': ['.number *= 2'], 'b': ['.number *= -1']},
    {'if': ['.date >= 2020-01-02'], 'a': ['.number *= .85'], 'b': ['.number *= .15']},
    # Metadata assignment
    {'if': ['a==1'], 'then': ['b=2', 'c=3']},
])
def test_valid_rewrite_rule(source):
    assert rewrite.RewriteRule(source)

@pytest.mark.parametrize('source', [
    # Incomplete rules
    {},
    {'if': ['.account in Equity']},
    {'a': ['.account = Income:Other'], 'b': ['.account = Expenses:Other']},
    # Condition/assignment mixup
    {'if': ['.account = Equity:Other'], 'then': ['equity-type = other']},
    {'if': ['.account == Equity:Other'], 'then': ['equity-type != other']},
    # Cross-category account assignment
    {'if': ['.date >= 2020-01-01'], 'then': ['.account = Assets:Cash']},
    {'if': ['.account in Equity'], 'then': ['.account = Assets:Cash']},
    # Number reallocation != 1
    {'if': ['.date >= 2020-01-01'], 'then': ['.number *= .5']},
    {'if': ['.date >= 2020-01-01'], 'a': ['k1=v1'], 'b': ['k2=v2']},
    # Date assignment
    {'if': ['.date == 2020-01-01'], 'then': ['.date = 2020-02-02']},
    # Redundant assignments
    {'if': ['.account in Income'],
     'then': ['.account = Income:Other', '.account = Income:Other']},
    {'if': ['.number > 0'],
     'a': ['.number *= .5', '.number *= .5'],
     'b': ['.number *= .5']},
])
def test_invalid_rewrite_rule(source):
    with pytest.raises(ValueError):
        rewrite.RewriteRule(source)

def test_rewrite_ruleset():
    account = 'Income:CurrencyConversion'
    ruleset = rewrite.RewriteRuleset(rewrite.RewriteRule(src) for src in [
        {'if': ['.account == Expenses:CurrencyConversion'],
         'rename': [f'.account = {account}']},
        {'if': ['project == alpha', '.account != Assets:Cash'],
         'cap': ['project = Alpha']},
    ])
    txn = testutil.Transaction(project='alpha', postings=[
        ('Assets:Cash', -20),
        ('Expenses:CurrencyConversion', 18),
        ('Expenses:CurrencyConversion', 1, {'project': 'Conservancy'}),
        ('Expenses:BankingFees', 1),
    ])
    posts = ruleset.rewrite(data.Posting.from_txn(txn))
    post = next(posts)
    assert post.account == 'Assets:Cash'
    assert post.meta['project'] == 'alpha'
    post = next(posts)
    assert post.account == account
    # Project not capitalized because the first rule took priority
    assert post.meta['project'] == 'alpha'
    post = next(posts)
    assert post.account == account
    assert post.meta['project'] == 'Conservancy'
    post = next(posts)
    assert post.account == 'Expenses:BankingFees'
    assert post.meta['project'] == 'Alpha'

def test_ruleset_from_yaml_path():
    yaml_path = testutil.test_path('userconfig/Rewrites01.yml')
    assert rewrite.RewriteRuleset.from_yaml(yaml_path)

def test_ruleset_from_yaml_str():
    with testutil.test_path('userconfig/Rewrites01.yml').open() as yaml_file:
        yaml_s = yaml_file.read()
    assert rewrite.RewriteRuleset.from_yaml(yaml_s)

def test_bad_ruleset_yaml_path():
    yaml_path = testutil.test_path('repository/Projects/project-data.yml')
    with pytest.raises(ValueError):
        rewrite.RewriteRuleset.from_yaml(yaml_path)

@pytest.mark.parametrize('source', [
    # Wrong root objects
    1,
    2.3,
    True,
    None,
    {},
    'string',
    [{}, 'a'],
    [{}, ['b']],
    # Rules have wrong type
    [{'if': '.account in Equity', 'add': ['testkey = value']}],
    [{'if': ['.account in Equity'], 'add': 'testkey = value'}],
])
def test_bad_ruleset_yaml_str(source):
    yaml_doc = yaml.safe_dump(source)
    with pytest.raises(ValueError):
        rewrite.RewriteRuleset.from_yaml(yaml_doc)