Define custom DoesNotExist for models.

4 files changed, 46 insertions, 27 deletions

diff --git a/talerbank/app/middleware.py b/talerbank/app/middleware.py
index 1af6f06..74f684e 100644
--- a/talerbank/app/middleware.py
+++ b/talerbank/app/middleware.py
@@ -2,7 +2,8 @@ import logging
 import zlib
 from django.http import JsonResponse
 from django.shortcuts import redirect
-from .models import BankAccount, BankTransaction
+from .models import BankAccount, BankTransaction, \
+    BankAccountDoesNotExist, BankTransactionDoesNotExist
 from .views import \
     (DebitLimitException, SameAccountException,
      LoginFailed, RejectNoRightsException)
@@ -70,7 +71,8 @@ class ExceptionMiddleware:
         # List of all the exceptions that are managed by
         # this module.
         self.excs = {
-            BankAccount.DoesNotExist: 0,
+            BankAccountDoesNotExist: 0,
+            BankTransactionDoesNotExist: 0,
             BankTransaction.DoesNotExist: 1,
             SameAccountException: 2,
             DebitLimitException: 3,
@@ -155,26 +157,10 @@ class ExceptionMiddleware:
         # page to be returned.
         render_to = self.render.get(request.path)
 
-        try:
-            hint = exception.hint
-            http_status_code = exception.http_status_code
-        ##
-        # This exception happens when the middleware is catching
-        # DoesNotExist exceptions; the ideal fix is to get BankAccount
-        # and BankTransaction classes to override their 'DoesNotExist'
-        # field with some custom class, but that wasn't straightforward
-        # (in the sense that on different systems we had different
-        # results, so we fallback on this more sound / manual approach)
-        except AttributeError:
-            hint = "The database (BankAccount / BankTransaction) object wasn't found."
-            http_status_code = 404
-
         if not render_to:
-            return JsonResponse({
-                "ec": taler_ec,
-                "error": hint
-            },
-                                status=http_status_code)
+            return JsonResponse({"ec": taler_ec,
+                                 "error": exception.hint},
+                                status=exception.http_status_code)
         request.session["profile_hint"] = \
             True, False, exception.hint
         return redirect(render_to)
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):
diff --git a/talerbank/app/tests.py b/talerbank/app/tests.py
index 2771bc7..7595bb0 100644
--- a/talerbank/app/tests.py
+++ b/talerbank/app/tests.py
@@ -30,7 +30,8 @@ from django.urls import reverse
 from django.conf import settings
 from django.contrib.auth.models import User
 from mock import patch, MagicMock
-from .models import BankAccount, BankTransaction
+from .models import BankAccount, BankTransaction, \
+    BankAccountDoesNotExist, BankTransactionDoesNotExist
 from . import urls
 from .views import wire_transfer
 from .amount import Amount, CurrencyMismatch, BadFormatAmount
@@ -542,6 +543,13 @@ class HistoryContext:
     def dump(self):
         return self.urlargs
 
+class CustomDoesNotExistTestCase(TestCase):
+
+    def test_bankaccount_doesnotexist(self):
+        with self.assertRaises(BankAccountDoesNotExist):
+            BankAccount.objects.get(account_no=1000)
+        with self.assertRaises(BankTransactionDoesNotExist):
+            BankTransaction.objects.get(subject="1000")
 
 class HistoryTestCase(TestCase):
     def setUp(self):
diff --git a/talerbank/app/views.py b/talerbank/app/views.py
index a5e7b26..67b1927 100644
--- a/talerbank/app/views.py
+++ b/talerbank/app/views.py
@@ -62,7 +62,6 @@ LOGGER = logging.getLogger(__name__)
 # can handle (because of the wallet).
 UINT64_MAX = (2**64) - 1
 
-
 ##
 # Exception raised upon failing login.
 #