From 3ffa66117ba2009bc1ba8af28cd3f18f5defc1aa Mon Sep 17 00:00:00 2001
From: Martijn Hemeryck <martijn.hemeryck@gmail.com>
Date: Fri, 21 May 2021 14:13:49 +0200
Subject: [PATCH] fixup! Implements BE BBAN checksum check

---
 schwifty/checksum/__init__.py | 1 +
 schwifty/checksum/belgium.py  | 8 ++++----
 schwifty/iban.py              | 3 +++
 tests/test_iban.py            | 1 +
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/schwifty/checksum/__init__.py b/schwifty/checksum/__init__.py
index 8c88029..29f4be9 100644
--- a/schwifty/checksum/__init__.py
+++ b/schwifty/checksum/__init__.py
@@ -37,3 +37,4 @@ def register(algorithm_cls: Type[Algorithm], prefix: Optional[str] = None) -> Ty
 
 from schwifty.checksum import germany  # noqa
 from schwifty.checksum import italy  # noqa
+from schwifty.checksum import belgium  # noqa
diff --git a/schwifty/checksum/belgium.py b/schwifty/checksum/belgium.py
index 2bfcc2a..6f1ab98 100644
--- a/schwifty/checksum/belgium.py
+++ b/schwifty/checksum/belgium.py
@@ -1,15 +1,15 @@
 import functools
 
-import schwifty.checksum
+from schwifty import checksum
 
 
-register = functools.partial(schwifty.checksum.register, prefix="BE")
+register = functools.partial(checksum.register, prefix="BE")
 
 
 @register
-class DefaultAlgorithm(schwifty.checksum.Algorithm):
+class DefaultAlgorithm(checksum.Algorithm):
     name = "default"
-    accepts = schwifty.checksum.InputType.BBAN
+    accepts = checksum.InputType.BBAN
 
     def compute(self, bban: str) -> str:
         return str(int(bban[:-2]) % 97)
diff --git a/schwifty/iban.py b/schwifty/iban.py
index eb014b2..35d6c93 100644
--- a/schwifty/iban.py
+++ b/schwifty/iban.py
@@ -43,6 +43,9 @@ def add_bban_checksum(country_code: str, bban: str) -> str:
     if country_code == "IT":
         checksum = algorithms["IT:default"].compute(bban[1:])
         bban = checksum + bban[1:]
+    elif country_code == "BE":
+        checksum = algorithms["BE:default"].compute(bban)
+        bban = bban[:-2] + checksum
     return bban
 
 
diff --git a/tests/test_iban.py b/tests/test_iban.py
index b0332c1..cecad39 100644
--- a/tests/test_iban.py
+++ b/tests/test_iban.py
@@ -154,6 +154,7 @@ def test_iban_properties():
         (("GB", "NWBK", "31926819", "601613"), "GB29NWBK60161331926819"),
         (("GB", "NWBK", "31926819"), "GB66NWBK00000031926819"),
         (("GB", "NWBK601613", "31926819"), "GB29NWBK60161331926819"),
+        (("BE", "050", "123456"), "BE67050000123487"),
     ],
 )
 def test_generate_iban(components, compact):