Issue #10; Local backend short url domain

3 files changed, 15 insertions, 2 deletions

diff --git a/brevisurl/backends/base.py b/brevisurl/backends/base.py
index 3fc053a..f667e79 100644
--- a/brevisurl/backends/base.py
+++ b/brevisurl/backends/base.py
@@ -10,6 +10,7 @@ class BaseBrevisUrlBackend(object):
     def __init__(self, fail_silently=False, **kwargs):
         self.fail_silently = fail_silently
         self.class_path = '{0}.{1}'.format(self.__module__, self.__class__.__name__)
+        self.kwargs = kwargs
 
     def open(self):
         """Open a network connection.
diff --git a/brevisurl/backends/local.py b/brevisurl/backends/local.py
index 9cb37a6..741904d 100644
--- a/brevisurl/backends/local.py
+++ b/brevisurl/backends/local.py
@@ -36,7 +36,11 @@ class BrevisUrlBackend(BaseBrevisUrlBackend):
 
         try:
             short_url = ShortUrl()
-            if brevisurl.settings.LOCAL_BACKEND_DOMAIN is not None:
+            if self.kwargs.get('domain') is not None:
+                short_url.shortened_url = '{0}{1}'.format(self.kwargs.get('domain').rstrip('/'),
+                                                          reverse('brevisurl_redirect',
+                                                                  kwargs={'token': self.__generate_token()}))
+            elif brevisurl.settings.LOCAL_BACKEND_DOMAIN is not None:
                 short_url.shortened_url = '{0}{1}'.format(brevisurl.settings.LOCAL_BACKEND_DOMAIN.rstrip('/'),
                                                           reverse('brevisurl_redirect',
                                                                   kwargs={'token': self.__generate_token()}))
diff --git a/brevisurl/tests/backends/test_local.py b/brevisurl/tests/backends/test_local.py
index 39ab8c1..620b9a4 100644
--- a/brevisurl/tests/backends/test_local.py
+++ b/brevisurl/tests/backends/test_local.py
@@ -87,4 +87,12 @@ class TestLocalBrevisUrlBackend(TestCase):
         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
+        brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS = _default_chars
+
+    def test_custom_domain(self):
+        original_url = 'http://www.codescale.net/'
+        connection = get_connection('brevisurl.backends.local.BrevisUrlBackend', domain='http://test.com/')
+        self.assertEqual(ShortUrl.objects.all().count(), 0)
+        short_url = connection.shorten_url(original_url)
+        self.assertEqual(ShortUrl.objects.all().count(), 1)
+        self.assertRegexpMatches(short_url.shortened_url, r'^http://test\.com/.{5}')
\ No newline at end of file