From 0413fed8b957bc543df0f044020176304660c5e5 2020-04-06 20:03:56
From: Brett Smith <brettcsmith@brettcsmith.org>
Date: 2020-04-06 20:03:56
Subject: [PATCH] meta_entity: Use payee as entity when metadata not available. RT#10529.

---

diff --git a/conservancy_beancount/data.py b/conservancy_beancount/data.py
index 1af27ac65b64cf649492e8d1a9ed246b2de6327a..8e191dfbf6677cc57602ab4f1d21d09fdc621fdf 100644
--- a/conservancy_beancount/data.py
+++ b/conservancy_beancount/data.py
@@ -185,6 +185,15 @@ class PostingMeta(Metadata):
         else:
             self.meta = collections.ChainMap(post.meta, txn.meta)
 
+    def __getitem__(self, key: MetaKey) -> MetaValue:
+        try:
+            return super().__getitem__(key)
+        except KeyError:
+            if key == 'entity' and self.txn.payee is not None:
+                return self.txn.payee
+            else:
+                raise
+
     def __setitem__(self, key: MetaKey, value: MetaValue) -> None:
         if self.post.meta is None:
             self.post = self.post._replace(meta={key: value})
diff --git a/conservancy_beancount/plugin/meta_entity.py b/conservancy_beancount/plugin/meta_entity.py
index 78d562f356e7d2f6fd4889e0ef0c1a408035faa3..7f4ecf1ffad82d85c51ffe46d3d4a8e0309975f8 100644
--- a/conservancy_beancount/plugin/meta_entity.py
+++ b/conservancy_beancount/plugin/meta_entity.py
@@ -49,7 +49,7 @@ class MetaEntity(core.TransactionHook):
     del alnum
 
     def run(self, txn: Transaction) -> errormod.Iter:
-        txn_entity = txn.meta.get(self.METADATA_KEY)
+        txn_entity = txn.meta.get(self.METADATA_KEY, txn.payee)
         if txn_entity is None:
             txn_entity_ok = None
         elif isinstance(txn_entity, str):
diff --git a/tests/test_data_posting_meta.py b/tests/test_data_posting_meta.py
index 1edfc66203e7723438096009b450bc8db4e1a777..0fb50026488064e06996cdcf25bf87296a93ffe2 100644
--- a/tests/test_data_posting_meta.py
+++ b/tests/test_data_posting_meta.py
@@ -21,13 +21,21 @@ from . import testutil
 from conservancy_beancount import data
 
 @pytest.fixture
-def simple_txn(index=None, key=None):
+def simple_txn():
     return testutil.Transaction(note='txn note', postings=[
         ('Assets:Cash', 5),
         ('Income:Donations', -5, {'note': 'donation love', 'extra': 'Extra'}),
     ])
 SIMPLE_TXN_METAKEYS = frozenset(['filename', 'lineno', 'note'])
 
+@pytest.fixture
+def payee_txn():
+    return testutil.Transaction(payee='SampleCo', postings=[
+        ('Assets:Receivable:Accounts', -100),
+        ('Assets:Checking', 95),
+        ('Expenses:BankingFees', 5, {'entity': 'MyBank'}),
+    ])
+
 def test_getitem_transaction(simple_txn):
     assert data.PostingMeta(simple_txn, 0)['note'] == 'txn note'
 
@@ -88,6 +96,22 @@ def test_get_links_from_post_override(simple_txn):
     meta = data.PostingMeta(simple_txn, 1)
     assert list(meta.get_links('note')) == ['donation', 'love']
 
+def test_payee_used_as_entity(payee_txn):
+    actual = [data.PostingMeta(payee_txn, n, p)['entity']
+              for n, p in enumerate(payee_txn.postings)]
+    assert actual == ['SampleCo', 'SampleCo', 'MyBank']
+
+def test_entity_metadata_has_precedence_over_payee(payee_txn):
+    payee_txn.meta['entity'] = 'ExampleCo'
+    actual = [data.PostingMeta(payee_txn, n, p)['entity']
+              for n, p in enumerate(payee_txn.postings)]
+    assert actual == ['ExampleCo', 'ExampleCo', 'MyBank']
+
+def test_keyerror_when_no_entity_or_payee(simple_txn):
+    meta = data.PostingMeta(simple_txn, 1)
+    with pytest.raises(KeyError):
+        meta['entity']
+
 # The .get() tests are arguably testing the stdlib, but they're short and
 # they confirm that we're using the stdlib as we intend.
 def test_get_with_meta_value(simple_txn):
diff --git a/tests/test_meta_entity.py b/tests/test_meta_entity.py
index 4542fa4769b22f9acb0296e98a58b92ef02a684d..f9caac672f38310b313b0a601b489658a7dbc7a2 100644
--- a/tests/test_meta_entity.py
+++ b/tests/test_meta_entity.py
@@ -110,6 +110,36 @@ def test_invalid_values_on_transactions(hook, src_value):
     assert all(error.message == "transaction has invalid entity: {}".format(src_value)
                for error in hook.run(txn))
 
+@pytest.mark.parametrize('src_value', VALID_VALUES)
+def test_valid_values_on_payee(hook, src_value):
+    txn = testutil.Transaction(payee=src_value, postings=[
+        ('Assets:Cash', -25),
+        ('Expenses:General', 25),
+    ])
+    assert not any(hook.run(txn))
+
+@pytest.mark.parametrize('src_value', INVALID_VALUES)
+def test_invalid_values_on_payee(hook, src_value):
+    txn = testutil.Transaction(payee=src_value, postings=[
+        ('Assets:Cash', -25),
+        ('Expenses:General', 25),
+    ])
+    errors = list(hook.run(txn))
+    assert 1 <= len(errors) <= 2
+    assert all(error.message == "transaction has invalid entity: {}".format(src_value)
+               for error in hook.run(txn))
+
+@pytest.mark.parametrize('payee,src_value', testutil.combine_values(
+    INVALID_VALUES,
+    VALID_VALUES,
+))
+def test_invalid_payee_but_valid_metadata(hook, payee, src_value):
+    txn = testutil.Transaction(**{'payee': payee, TEST_KEY: src_value}, postings=[
+        ('Assets:Cash', -25),
+        ('Expenses:Other', 25),
+    ])
+    assert not any(hook.run(txn))
+
 @pytest.mark.parametrize('account,required', [
     ('Assets:Bank:Checking', False),
     ('Assets:Cash', False),