[Refactoring] Split extractor from Nordigen Client

pfbudget/core/manager.py

@@ -28,9 +28,9 @@ from pfbudget.db.model import (
     TransactionCategory,
 )
 from pfbudget.extract.credentials import Credentials
-from pfbudget.extract.extract import Extract
+from pfbudget.extract.nordigen import NordigenClient
-from pfbudget.extract.psd2 import PSD2Client
 from pfbudget.extract.parsers import parse_data
+from pfbudget.extract.psd2 import PSD2Extractor
 
 dotenv.load_dotenv()
 
@@ -85,16 +85,18 @@ class Manager:
                     else:
                         banks = session.get(Bank, Bank.name, params[3])
                     session.expunge_all()
-                client.start = params[0]
+
-                client.end = params[1]
+                extractor = PSD2Extractor(client)
-                transactions = client.extract(banks)
+                transactions = []
+                for bank in banks:
+                    transactions.extend(extractor.extract(bank, params[0], params[1]))
 
                 # dry-run
                 if not params[2]:
                     with self.db.session() as session:
                         session.add(sorted(transactions))
                 else:
-                    print(transactions)
+                    print(sorted(transactions))
 
             case Operation.Categorize:
                 with self.db.session() as session:
@@ -419,10 +421,10 @@ class Manager:
         self._db = url
 
     @staticmethod
-    def nordigen_client() -> Extract:
+    def nordigen_client() -> NordigenClient:
         credentials = Credentials(
             os.environ.get("SECRET_ID"),
             os.environ.get("SECRET_KEY"),
             os.environ.get("TOKEN"),
         )
-        return PSD2Client(credentials)
+        return NordigenClient(credentials)

pfbudget/extract/exceptions.py

@@ -6,5 +6,13 @@ class BankError(ExtractError):
     pass
 
 
-class CredentialsError(ExtractError):
+class PSD2ClientError(ExtractError):
+    pass
+
+
+class CredentialsError(PSD2ClientError):
+    pass
+
+
+class DownloadError(PSD2ClientError):
     pass

pfbudget/extract/nordigen.py

@@ -0,0 +1,92 @@
+import json
+import nordigen
+import requests
+import time
+import uuid
+
+from .credentials import Credentials
+from .exceptions import CredentialsError, DownloadError
+
+
+class NordigenClient:
+    redirect_url = "https://murta.dev"
+
+    def __init__(self, credentials: Credentials):
+        super().__init__()
+
+        if not credentials.valid():
+            raise CredentialsError
+
+        self._client = nordigen.NordigenClient(
+            secret_key=credentials.key, secret_id=credentials.id, timeout=5
+        )
+
+        if credentials.token:
+            self._client.token = credentials.token
+
+    def download(self, requisition_id):
+        try:
+            requisition = self._client.requisition.get_requisition_by_id(requisition_id)
+            print(requisition)
+        except requests.HTTPError as e:
+            raise DownloadError(e)
+
+        transactions = {}
+        for acc in requisition["accounts"]:
+            account = self._client.account_api(acc)
+
+            retries = 0
+            while retries < 3:
+                try:
+                    downloaded = account.get_transactions()
+                    break
+                except requests.ReadTimeout:
+                    retries += 1
+                    print(f"Request #{retries} timed-out, retrying in 1s")
+                    time.sleep(1)
+
+            if not downloaded:
+                print(f"Couldn't download transactions for {account}")
+                continue
+
+            transactions.update(downloaded)
+
+        return transactions
+
+    def dump(self, bank, downloaded):
+        with open("json/" + bank.name + ".json", "w") as f:
+            json.dump(downloaded, f)
+
+    def generate_token(self):
+        self.token = self._client.generate_token()
+        print(f"New access token: {self.token}")
+        return self.token
+
+    def requisition(self, id: str, country: str = "PT"):
+        requisition = self._client.initialize_session(
+            redirect_uri=self.redirect_url,
+            institution_id=id,
+            reference_id=str(uuid.uuid4()),
+        )
+        return requisition.link, requisition.requisition_id
+
+    def country_banks(self, country: str):
+        return self._client.institution.get_institutions(country)
+
+    # def __token(self):
+    #     if token := os.environ.get("TOKEN"):
+    #         return token
+    #     else:
+    #         token = self._client.generate_token()
+    #         print(f"New access token: {token}")
+    #         return token["access"]
+
+    @property
+    def token(self):
+        return self._token
+
+    @token.setter
+    def token(self, value):
+        if self._token:
+            print("Replacing existing token with {value}")
+        self._token = value

pfbudget/extract/psd2.py

@@ -1,142 +1,39 @@
 import datetime as dt
-import json
-import nordigen
-import requests
-import time
-import uuid
 
 from typing import Sequence
 
-import pfbudget.db.model as t
+from pfbudget.db.model import Bank, BankTransaction
 from pfbudget.utils.converters import convert
 
-from .credentials import Credentials
+from .exceptions import BankError, DownloadError, ExtractError
-from .exceptions import BankError, CredentialsError, ExtractError
 from .extract import Extract
+from .nordigen import NordigenClient
 
 
-class PSD2Client(Extract):
+class PSD2Extractor(Extract):
-    redirect_url = "https://murta.dev"
+    def __init__(self, client: NordigenClient):
+        self.__client = client
 
-    def __init__(self, credentials: Credentials):
+    def extract(
-        super().__init__()
+        self, bank: Bank, start=dt.date.min, end=dt.date.max
+    ) -> Sequence[BankTransaction]:
+        if not bank.nordigen:
+            raise BankError("Bank doesn't have Nordigen info")
 
-        if not credentials.valid():
-            raise CredentialsError
-
-        self._client = nordigen.NordigenClient(
-            secret_key=credentials.key, secret_id=credentials.id, timeout=5
-        )
-
-        if credentials.token:
-            self._client.token = credentials.token
-
-        self._start = dt.date.min
-        self._end = dt.date.max
-
-    def extract(self, banks: Sequence[t.Bank]) -> list[t.BankTransaction]:
-        transactions = []
-        if not banks or any(not b.nordigen for b in banks):
-            raise BankError
-
-        for bank in banks:
-            downloaded = None
         try:
             print(f"Downloading from {bank}...")
-                downloaded = self.download(bank.nordigen.requisition_id)
+            downloaded = self.__client.download(bank.nordigen.requisition_id)
-            except requests.HTTPError as e:
+        except DownloadError as e:
-                print(f"There was an issue downloading from {bank.name} -> {e}")
+            print(f"There was an issue downloading from {bank.name}\n{e}")
             raise ExtractError(e)
 
-            if downloaded:
+        self.__client.dump(bank, downloaded)
-                self.dump(bank, downloaded)
 
-                converted = [
+        return [
-                    convert(t, bank) for t in downloaded["transactions"]["booked"]
+            t
+            for t in self.convert(bank, downloaded, start, end)
+            if start <= t.date <= end
         ]
 
-                transactions.extend(
+    def convert(self, bank, downloaded, start, end):
-                    [t for t in converted if self._start <= t.date <= self._end]
+        return [convert(t, bank) for t in downloaded["transactions"]["booked"]]
-                )
-
-        return sorted(transactions)
-
-    def download(self, requisition_id):
-        requisition = self._client.requisition.get_requisition_by_id(requisition_id)
-        print(requisition)
-
-        transactions = {}
-        for acc in requisition["accounts"]:
-            account = self._client.account_api(acc)
-
-            retries = 0
-            while retries < 3:
-                try:
-                    downloaded = account.get_transactions()
-                    break
-                except requests.ReadTimeout:
-                    retries += 1
-                    print(f"Request #{retries} timed-out, retrying in 1s")
-                    time.sleep(1)
-
-            if not downloaded:
-                print(f"Couldn't download transactions for {account}")
-                continue
-
-            transactions.update(downloaded)
-
-        return transactions
-
-    def dump(self, bank, downloaded):
-        with open("json/" + bank.name + ".json", "w") as f:
-            json.dump(downloaded, f)
-
-    def generate_token(self):
-        self.token = self._client.generate_token()
-        print(f"New access token: {self.token}")
-        return self.token
-
-    def requisition(self, id: str, country: str = "PT"):
-        requisition = self._client.initialize_session(
-            redirect_uri=self.redirect_url,
-            institution_id=id,
-            reference_id=str(uuid.uuid4()),
-        )
-        return requisition.link, requisition.requisition_id
-
-    def country_banks(self, country: str):
-        return self._client.institution.get_institutions(country)
-
-    @property
-    def start(self):
-        return self._start
-
-    @start.setter
-    def start(self, value):
-        self._start = value
-
-    @property
-    def end(self):
-        return self._end
-
-    @end.setter
-    def end(self, value):
-        self._end = value
-
-    # def __token(self):
-    #     if token := os.environ.get("TOKEN"):
-    #         return token
-    #     else:
-    #         token = self._client.generate_token()
-    #         print(f"New access token: {token}")
-    #         return token["access"]
-
-    @property
-    def token(self):
-        return self._token
-
-    @token.setter
-    def token(self, value):
-        if self._token:
-            print("Replacing existing token with {value}")
-        self._token = value

tests/test_psd2.py

@@ -8,7 +8,8 @@ import mocks.nordigen as mock
 from pfbudget.db.model import Bank, BankTransaction, Nordigen
 from pfbudget.extract.credentials import Credentials
 from pfbudget.extract.exceptions import BankError, CredentialsError
-from pfbudget.extract.psd2 import PSD2Client
+from pfbudget.extract.nordigen import NordigenClient
+from pfbudget.extract.psd2 import PSD2Extractor
 
 
 class MockGet:
@@ -55,41 +56,37 @@ def mock_requests(monkeypatch):
 
 
 @pytest.fixture
-def client() -> PSD2Client:
+def extractor() -> NordigenClient:
     credentials = Credentials("ID", "KEY", "TOKEN")
-    return PSD2Client(credentials)
+    return PSD2Extractor(NordigenClient(credentials))
 
 
 @pytest.fixture
-def banks() -> list[Bank]:
+def bank() -> list[Bank]:
     bank = Bank("Bank#1", "", "")
     bank.nordigen = Nordigen("", "", mock.id, False)
-    return [bank]
+    return bank
 
 
 class TestExtractPSD2:
     def test_empty_credentials(self):
         cred = Credentials("", "")
         with pytest.raises(CredentialsError):
-            PSD2Client(cred)
+            NordigenClient(cred)
 
-    def test_empty_banks(self, client):
+    def test_no_psd2_bank(self, extractor):
         with pytest.raises(BankError):
-            client.extract([])
+            extractor.extract(Bank("", "", ""))
 
-    def test_no_psd2_bank(self, client):
+    def test_timeout(self, monkeypatch, extractor, bank):
-        with pytest.raises(BankError):
-            client.extract([Bank("", "", "")])
-
-    def test_timeout(self, monkeypatch, client, banks):
         monkeypatch.setattr(
             "requests.get", MockGet(mock_exception=requests.ReadTimeout)
         )
         with pytest.raises(requests.Timeout):
-            client.extract(banks)
+            extractor.extract(bank)
 
-    def test_extract(self, client, banks):
+    def test_extract(self, extractor, bank):
-        assert client.extract(banks) == [
+        assert extractor.extract(bank) == [
             BankTransaction(
                 dt.date(2023, 1, 14), "string", Decimal("328.18"), "Bank#1"
             ),