Define custom DoesNotExist for models.

1 files changed, 30 insertions, 4 deletions

diff --git a/talerbank/app/models.py b/talerbank/app/models.py
index aa5d300..a5e1b41 100644
--- a/talerbank/app/models.py
+++ b/talerbank/app/models.py
@@ -28,7 +28,6 @@ from django.core.exceptions import \
     ObjectDoesNotExist
 from .amount import Amount, BadFormatAmount, NumberTooBig, CurrencyMismatch
 
-
 class InvalidAmount(Amount):
     def __init__(self, currency):
         super(InvalidAmount, self
@@ -121,9 +120,35 @@ class AmountField(models.Field):
             return Amount.parse(value)
         except BadFormatAmount:
             raise ValidationError(
-                "Invalid input for an amount string: %s" % value
-            )
+                "Invalid input for an amount string: %s" % value)
+
+class BankAccountDoesNotExist(Exception):
+    def __init__(self):
+        self.hint = "Bank account not found"
+        self.http_status_code = 404
+        self.minor_error_code = 0
+
+class BankTransactionDoesNotExist(Exception):
+    def __init__(self):
+        self.hint = "Bank transaction not found"
+        self.http_status_code = 404
+        self.minor_error_code = 0
+
+class CustomManager(models.Manager):
 
+    def __init__(self):
+        super(CustomManager, self).__init__()
+
+    def get_queryset(self):
+        return models.QuerySet(self.model, using=self._db)
+
+    def get(self, *args, **kwargs):
+        try:
+            return super(CustomManager, self).get(*args, **kwargs)
+        except BankAccount.DoesNotExist:
+            raise BankAccountDoesNotExist()
+        except BankTransaction.DoesNotExist:
+            raise BankTransactionDoesNotExist()
 
 ##
 # The class representing a bank account.
@@ -133,7 +158,7 @@ class BankAccount(models.Model):
     account_no = models.AutoField(primary_key=True)
     user = models.OneToOneField(User, on_delete=models.CASCADE)
     amount = AmountField(default=get_zero_amount)
-
+    objects = CustomManager()
 
 ##
 # The class representing a bank transaction.
@@ -154,6 +179,7 @@ class BankTransaction(models.Model):
     subject = models.CharField(default="(no subject given)", max_length=200)
     date = models.DateTimeField(auto_now=True, db_index=True)
     cancelled = models.BooleanField(default=False)
+    objects = CustomManager()
 
 
 class TalerWithdrawOperation(models.Model):