Support custom routes in Router

2 files changed, 79 insertions, 5 deletions

diff --git a/sugar_network/toolkit/router.py b/sugar_network/toolkit/router.py
index d6d34f6..f659a68 100644
--- a/sugar_network/toolkit/router.py
+++ b/sugar_network/toolkit/router.py
@@ -50,6 +50,18 @@ class Unauthorized(HTTPStatus):
     headers = {'WWW-Authenticate': 'Sugar'}
 
 
+def route(method, path):
+    path = path.strip('/').split('/')
+    # Only top level paths for now
+    enforce(len(path) == 1)
+
+    def decorate(func):
+        func.route = (method, path[0])
+        return func
+
+    return decorate
+
+
 class Router(object):
 
     def __init__(self, cp):
@@ -58,6 +70,16 @@ class Router(object):
         self._valid_origins = set()
         self._invalid_origins = set()
         self._host = None
+        self._routes = {}
+
+        cls = self.__class__
+        while cls is not None:
+            for name in dir(cls):
+                attr = getattr(self, name)
+                if hasattr(attr, 'route'):
+                    self._routes[attr.route] = attr
+            # pylint: disable-msg=E1101
+            cls = cls.__base__
 
         if 'SSH_ASKPASS' in os.environ:
             # Otherwise ssh-keygen will popup auth dialogs on registeration
@@ -105,7 +127,13 @@ class Router(object):
             enforce(isfile(static_path), 'No such file')
             result = file(static_path)
         else:
-            result = self._cp.call(request, response)
+            rout = None
+            if request.path:
+                rout = self._routes.get((request['method'], request.path[0]))
+            if rout:
+                result = rout(request, response)
+            else:
+                result = self._cp.call(request, response)
 
         if hasattr(result, 'read'):
             # pylint: disable-msg=E1103
@@ -136,8 +164,7 @@ class Router(object):
             response['Location'] = error.location
             response.content_type = None
         except Exception, error:
-            util.exception('Error while processing %r request',
-                    environ['PATH_INFO'] or '/')
+            util.exception('Error while processing %r request', request.url)
 
             if isinstance(error, ad.NotFound):
                 response.status = '404 Not Found'
@@ -152,7 +179,7 @@ class Router(object):
 
             if result is None:
                 result = {'error': str(error),
-                          'request': environ['PATH_INFO'] or '/',
+                          'request': request.url,
                           }
                 response.content_type = 'application/json'
 
diff --git a/tests/units/router.py b/tests/units/router.py
index e9908c2..7b9c5ca 100755
--- a/tests/units/router.py
+++ b/tests/units/router.py
@@ -13,7 +13,7 @@ from __init__ import tests
 
 import active_document as ad
 from sugar_network import node
-from sugar_network.toolkit.router import Router, _Request, _parse_accept_language, Unauthorized
+from sugar_network.toolkit.router import Router, _Request, _parse_accept_language, Unauthorized, route
 from active_toolkit import util
 from sugar_network.resources.volume import Volume
 
@@ -268,6 +268,53 @@ class RouterTest(tests.Test):
                 ['ru', 'en', 'es'],
                 _parse_accept_language('ru;q=1,en;q=1,es;q=0.5'))
 
+    def test_CustomRoutes(self):
+        calls = []
+
+        class TestRouterBase(Router):
+
+            @route('GET', '/foo')
+            def route1(self, request, response):
+                calls.append('route1')
+
+        class TestRouter(TestRouterBase):
+
+            @route('PUT', '/foo')
+            def route2(self, request, response):
+                calls.append('route2')
+
+            @route('GET', '/bar')
+            def route3(self, request, response):
+                calls.append('route3')
+
+        class CommandsProcessor(object):
+
+            def call(self, request, response):
+                calls.append('default')
+
+        cp = CommandsProcessor()
+        router = TestRouter(cp)
+
+        [i for i in router({'PATH_INFO': '/', 'REQUEST_METHOD': 'GET'}, lambda *args: None)]
+        self.assertEqual(['default'], calls)
+        del calls[:]
+
+        [i for i in router({'PATH_INFO': '//foo//', 'REQUEST_METHOD': 'GET'}, lambda *args: None)]
+        self.assertEqual(['route1'], calls)
+        del calls[:]
+
+        [i for i in router({'PATH_INFO': '/foo', 'REQUEST_METHOD': 'PUT'}, lambda *args: None)]
+        self.assertEqual(['route2'], calls)
+        del calls[:]
+
+        [i for i in router({'PATH_INFO': '/foo', 'REQUEST_METHOD': 'POST'}, lambda *args: None)]
+        self.assertEqual(['default'], calls)
+        del calls[:]
+
+        [i for i in router({'PATH_INFO': '/bar/foo/probe', 'REQUEST_METHOD': 'GET'}, lambda *args: None)]
+        self.assertEqual(['route3'], calls)
+        del calls[:]
+
 
 class Document(ad.Document):