Cleanup after merging pull request from unaizalakain

3 files changed, 37 insertions, 8 deletions

diff --git a/brevisurl/__init__.py b/brevisurl/__init__.py
index 66f2efa..b929f1f 100644
--- a/brevisurl/__init__.py
+++ b/brevisurl/__init__.py
@@ -1,3 +1,9 @@
+import traceback
+try:
+    from cStringIO import StringIO
+except ImportError:
+    from StringIO import StringIO
+
 import brevisurl.settings
 from brevisurl.utils import load_object
 
@@ -34,4 +40,17 @@ def shorten_url(original_url, fail_silently=False, connection=None):
 
     """
     connection = connection or get_connection(fail_silently=fail_silently)
-    return connection.shorten_url(original_url)
\ No newline at end of file
+    return connection.shorten_url(original_url)
+
+
+class Error(Exception):
+    """Base django-brevisurl Error."""
+
+    def __init__(self, value):
+        s = StringIO()
+        traceback.print_exc(file=s)
+        self.value = (value, s.getvalue())
+        s.close()
+
+    def __str__(self):
+        return repr(self.value)
\ No newline at end of file
diff --git a/brevisurl/backends/local.py b/brevisurl/backends/local.py
index 69944e8..4dc81b8 100644
--- a/brevisurl/backends/local.py
+++ b/brevisurl/backends/local.py
@@ -1,3 +1,4 @@
+import math
 import string
 import random
 import logging
@@ -6,14 +7,17 @@ from django.contrib.sites.models import Site
 from django.core.urlresolvers import reverse
 
 import brevisurl.settings
+from brevisurl import Error
 from brevisurl.backends.base import BaseBrevisUrlBackend
 from brevisurl.models import ShortUrl
 
 
 log = logging.getLogger(__name__)
 
-class TokensExhausted(Exception):
-    pass
+
+class TokensExhaustedError(Error):
+    """Exception is raised when tokens are exhausted."""
+
 
 class BrevisUrlBackend(BaseBrevisUrlBackend):
 
@@ -22,6 +26,7 @@ class BrevisUrlBackend(BaseBrevisUrlBackend):
     def shorten_url(self, original_url):
         """
         :raises: ImproperlyConfigured, django.core.exceptions.ValidationError
+        :raises: brevisurl.backends.local.TokensExhaustedError
         """
         try:
             short_url = ShortUrl.objects.get(backend=self.class_path, original_url=original_url)
@@ -55,11 +60,11 @@ class BrevisUrlBackend(BaseBrevisUrlBackend):
 
     def __generate_token(self, size=5):
         chars = list(string.ascii_letters + string.digits)
-        if ShortUrl.objects.count() >= len(chars) ** size:
-            raise TokensExhausted('Consider incrementing the token length or chars list')
+        if ShortUrl.objects.count() >= math.pow(len(chars), size):
+            raise TokensExhaustedError('Consider incrementing the token length or change the char list')
         random.shuffle(chars)
         while True:
             token = ''.join([random.choice(chars) for x in range(size)])
             if not ShortUrl.objects.filter(backend=self.class_path, shortened_url__endswith=token).count():
                 break
-        return token
+        return token
\ No newline at end of file
diff --git a/brevisurl/models.py b/brevisurl/models.py
index 0a583ad..9ade639 100644
--- a/brevisurl/models.py
+++ b/brevisurl/models.py
@@ -13,7 +13,7 @@ log = logging.getLogger(__name__)
 
 class ShortUrl(models.Model):
     """Model that represents shortened url."""
-    original_url = models.CharField(max_length=200, null=False, blank=False)
+    original_url = models.URLField(max_length=200, null=False, blank=False)
     original_url_hash = models.CharField(max_length=64, null=False, blank=False)
     shortened_url = models.URLField(max_length=200, null=False, blank=False, unique=True)
     backend = models.CharField(max_length=200, null=False, blank=False)
@@ -33,9 +33,14 @@ class ShortUrl(models.Model):
     def clean(self):
         url_validator = URLValidator()
         try:
+            url_validator(self.original_url)
+        except ValidationError:
+            log.exception('ShortUrl.original_url "%s" is not valid URL', self.original_url)
+            raise
+        try:
             url_validator(self.shortened_url)
         except ValidationError:
-            log.exception('ShortUrl.shortened_url is not valid URL')
+            log.exception('ShortUrl.shortened_url "%s" is not valid URL', self.shortened_url)
             raise
         return super(ShortUrl, self).clean()