Issue #7; configurable token characters

3 files changed, 36 insertions, 4 deletions

diff --git a/brevisurl/backends/local.py b/brevisurl/backends/local.py
index 4dc81b8..6c3f6b4 100644
--- a/brevisurl/backends/local.py
+++ b/brevisurl/backends/local.py
@@ -1,5 +1,4 @@
 import math
-import string
 import random
 import logging
 
@@ -59,7 +58,7 @@ class BrevisUrlBackend(BaseBrevisUrlBackend):
                 raise
 
     def __generate_token(self, size=5):
-        chars = list(string.ascii_letters + string.digits)
+        chars = brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS
         if ShortUrl.objects.count() >= math.pow(len(chars), size):
             raise TokensExhaustedError('Consider incrementing the token length or change the char list')
         random.shuffle(chars)
diff --git a/brevisurl/settings.py b/brevisurl/settings.py
index 2d54b04..32d7fbf 100644
--- a/brevisurl/settings.py
+++ b/brevisurl/settings.py
@@ -1,4 +1,13 @@
+import string
+
 from django.conf import settings
 
+
+# Default backend used to for get_connection() fuction
 DEFAULT_BACKEND = getattr(settings, 'BREVISURL_BACKEND', 'brevisurl.backends.local.BrevisUrlBackend')
-LOCAL_BACKEND_DOMAIN = getattr(settings, 'BREVISURL_BACKEND_LOCAL_DOMAIN', None)
\ No newline at end of file
+
+# Domain that is used to create shortened urls
+LOCAL_BACKEND_DOMAIN = getattr(settings, 'BREVISURL_BACKEND_LOCAL_DOMAIN', None)
+
+# Characters that are used to generate tokens for local backend
+LOCAL_BACKEND_TOKEN_CHARS = getattr(settings, 'BREVISURL_LOCAL_BACKEND_TOKEN_CHARS', list(string.ascii_letters + string.digits))
\ No newline at end of file
diff --git a/brevisurl/tests/backends/test_local.py b/brevisurl/tests/backends/test_local.py
index 86565d2..39ab8c1 100644
--- a/brevisurl/tests/backends/test_local.py
+++ b/brevisurl/tests/backends/test_local.py
@@ -5,6 +5,7 @@ from django.core.validators import URLValidator
 import brevisurl.settings
 from brevisurl import get_connection
 from brevisurl.models import ShortUrl
+from brevisurl.backends.local import TokensExhaustedError
 
 
 class TestLocalBrevisUrlBackend(TestCase):
@@ -63,4 +64,27 @@ class TestLocalBrevisUrlBackend(TestCase):
     def test_shorten_url_invalid_original_url_fail_silently(self):
         self.connection.fail_silently = True
         shorl_url = self.connection.shorten_url('www.codescale.')
-        self.assertIsNone(shorl_url)
\ No newline at end of file
+        self.assertIsNone(shorl_url)
+
+    def test_configurable_token_chars(self):
+        original_url = 'http://www.codescale.net/'
+        _default_chars = brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS
+        brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS = ['a']
+        self.assertEqual(ShortUrl.objects.all().count(), 0)
+        short_url = self.connection.shorten_url(original_url)
+        self.assertEqual(ShortUrl.objects.all().count(), 1)
+        self.assertEqual(short_url.original_url, original_url)
+        self.assertRegexpMatches(short_url.shortened_url, r'/aaaaa$')
+        brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS = _default_chars
+
+    def test_exhausted_tokens(self):
+        original_url = 'http://www.codescale.net/'
+        _default_chars = brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS
+        brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS = ['a']
+        self.assertEqual(ShortUrl.objects.all().count(), 0)
+        short_url = self.connection.shorten_url(original_url)
+        self.assertEqual(ShortUrl.objects.all().count(), 1)
+        with self.assertRaises(TokensExhaustedError):
+            original_url = 'http://www.codescale.net/another/'
+            short_url = self.connection.shorten_url(original_url)
+        brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS = _default_chars
\ No newline at end of file