Add test for external posting_predictors

jbms · Dec 29, 2020 · 50c6ba9 · 50c6ba9
1 parent 220ded3
commit 50c6ba9
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 2 deletions.
diff --git a/beancount_import/source/generic_importer_source_test.py b/beancount_import/source/generic_importer_source_test.py
@@ -4,6 +4,7 @@
 
 from .source_test import check_source_example
 from beancount.ingest.importers.csv import Importer as CSVImporter, Col
+from beancount.core.data import Transaction, Posting, Amount
 
 testdata_dir = os.path.realpath(
     os.path.join(
@@ -14,10 +15,14 @@
 examples = [
     'test_basic',
     'test_invalid',
-    'test_training_examples'
+    'test_training_examples',
+    'test_dummy_predictor'
 ]
+# tests that should be wrapped in below DummyPostingPredictor
+wrap_dummy_predictor = ['test_dummy_predictor']
 
-importer = CSVImporter({Col.DATE: 'Date',
+def Importer():
+    return CSVImporter({Col.DATE: 'Date',
                         Col.NARRATION1: 'Description',
                         Col.AMOUNT: 'Amount',
                         },
@@ -26,9 +31,39 @@
                        '"Date","Description","Amount"',
                        )
 
+class DummyPostingPredictor:
+    DUMMY_ACCOUNT = "Assets:Dummy"
+
+    def __init__(self, importer):
+        self.importer = importer
+        # move extract to _extract
+        self.importer._extract = self.importer.extract
+        self.importer.extract = self.extract_wrapper
+
+    def extract_wrapper(self, f, existing_entries):
+        entries = self.importer._extract(f, existing_entries)
+        for entry in entries:
+            if isinstance(entry, Transaction):
+                p = entry.postings[0]
+                entry.postings.append(
+                    Posting(
+                        account=self.DUMMY_ACCOUNT,
+                        units=Amount(currency=p.units.currency, number=-1*p.units.number),
+                        cost=None,
+                        price=None,
+                        flag=None,
+                        meta={},
+                    )
+                )
+        return entries
+
 
 @pytest.mark.parametrize('name', examples)
 def test_source(name: str):
+    importer = Importer()
+    # wrap the importer in dummy posting predictor
+    if name in wrap_dummy_predictor:
+        _ = DummyPostingPredictor(importer)
     check_source_example(
         example_dir=os.path.join(testdata_dir, name),
         source_spec={

diff --git a/testdata/source/generic_importer/test_dummy_predictor/accounts.txt b/testdata/source/generic_importer/test_dummy_predictor/accounts.txt
@@ -0,0 +1 @@
+Assets:Bank
diff --git a/testdata/source/generic_importer/test_dummy_predictor/import_results.beancount b/testdata/source/generic_importer/test_dummy_predictor/import_results.beancount
@@ -0,0 +1,65 @@
+;; date: 2020-01-01
+;; info: {"filename": "<testdata>/csv/generic_statement.csv", "line": 1, "type": "text/csv"}
+
+; features: []
+2020-01-01 * "by debit card-OTHPG 063441 GOOGLE CLOUD INDIA PVTTHANE-"
+  Assets:Bank   -1 USD
+    date: 2020-01-01
+    source_desc: "by debit card-OTHPG 063441 GOOGLE CLOUD INDIA PVTTHANE-"
+  Assets:Dummy   1 USD
+    unconfirmed_account: TRUE
+
+;; date: 2020-01-01
+;; info: {"filename": "<testdata>/csv/generic_statement.csv", "line": 2, "type": "text/csv"}
+
+; features: []
+2020-01-01 * "by debit card-OTHPG 063444 GOOGLE CLOUD INDIA PVTTHANE-"
+  Assets:Bank   -1 USD
+    date: 2020-01-01
+    source_desc: "by debit card-OTHPG 063444 GOOGLE CLOUD INDIA PVTTHANE-"
+  Assets:Dummy   1 USD
+    unconfirmed_account: TRUE
+
+;; date: 2020-01-02
+;; info: {"filename": "<testdata>/csv/generic_statement.csv", "line": 3, "type": "text/csv"}
+
+; features: []
+2020-01-02 * "BULK POSTING- 00000008237 250120 GOOGLE CLOUD INDIA PVT-"
+  Assets:Bank    1 USD
+    date: 2020-01-02
+    source_desc: "BULK POSTING- 00000008237 250120 GOOGLE CLOUD INDIA PVT-"
+  Assets:Dummy  -1 USD
+    unconfirmed_account: TRUE
+
+;; date: 2020-01-02
+;; info: {"filename": "<testdata>/csv/generic_statement.csv", "line": 5, "type": "text/csv"}
+
+; features: []
+2020-01-02 * "BULK POSTING- 00000008237 250120 GOOGLE CLOUD INDIA PVT-"
+  Assets:Bank    1 USD
+    date: 2020-01-02
+    source_desc: "BULK POSTING- 00000008237 250120 GOOGLE CLOUD INDIA PVT-"
+  Assets:Dummy  -1 USD
+    unconfirmed_account: TRUE
+
+;; date: 2020-01-02
+;; info: {"filename": "<testdata>/csv/generic_statement.csv", "line": 4, "type": "text/csv"}
+
+; features: []
+2020-01-02 * "ATM-WD Some Random ATM Machine"
+  Assets:Bank    500 USD
+    date: 2020-01-02
+    source_desc: "ATM-WD Some Random ATM Machine"
+  Assets:Dummy  -500 USD
+    unconfirmed_account: TRUE
+
+;; date: 2020-01-05
+;; info: {"filename": "<testdata>/csv/generic_statement.csv", "line": 6, "type": "text/csv"}
+
+; features: []
+2020-01-05 * "Transfer to 1234567890123"
+  Assets:Bank    300 USD
+    date: 2020-01-05
+    source_desc: "Transfer to 1234567890123"
+  Assets:Dummy  -300 USD
+    unconfirmed_account: TRUE
diff --git a/testdata/source/generic_importer/test_dummy_predictor/journal.beancount b/testdata/source/generic_importer/test_dummy_predictor/journal.beancount
@@ -0,0 +1 @@
+1900-01-01 open Assets:Bank
diff --git a/testdata/source/generic_importer/test_dummy_predictor/training_examples.json b/testdata/source/generic_importer/test_dummy_predictor/training_examples.json
@@ -0,0 +1 @@
+[]