better webclient.py

VERSION

@@ -1 +1 @@
-Version 2.00.1 (2012-08-29 18:31:20) rc4
+Version 2.00.1 (2012-08-29 21:57:05) rc4

gluon/contrib/webclient.py

@@ -2,8 +2,18 @@
 Developed by Massimo Di Pierro
 Released under the web2py license (LGPL)
 
-It an interface on top of urllib2 that allows authentication and understand
+It an interface on top of urllib2 which simplifies scripting of http requests
-web2py cookies and web2py forms. An example of usage is at the bottom.
+mostly for testing purposes
+
+- customizable
+- supports basic auth
+- supports cookies 
+- supports session cookies (tested with web2py sessions)
+- detects broken session
+- detects web2py form postbacks and handles formname and formkey
+- detects web2py tickets
+
+Some examples at the bottom.
 """
 
 import re
@@ -11,82 +21,153 @@ import time
 import urllib
 import urllib2
 
+DEFAULT_HEADERS = {
+    'user-agent': 'Mozilla/4.0', # some servers are picky
+    'accept-language': 'en',
+    }
+
+FORM_REGEX = re.compile('(\<input name\="_formkey" type\="hidden" value\="(?P<formkey>.+?)" \/\>)?\<input name\="_formname" type\="hidden" value\="(?P<formname>.+?)" \/\>')
+
+SESSION_REGEX = 'session_id_(?P<name>.+)'
 
 class WebClient(object):
-    regex = re.compile('\<input name\="_formkey" type\="hidden" value\="(?P<formkey>.+?)" \/\>\<input name\="_formname" type\="hidden" value\="(?P<formname>.+?)" \/\>')
 
-    def __init__(self,app='', postbacks=True):
+    def __init__(self,
+                 app = '',
+                 postbacks = True,
+                 default_headers = DEFAULT_HEADERS,
+                 session_regex = SESSION_REGEX):
+        self.app = app
         self.postbacks = postbacks
         self.history = []
-        self.app = app
         self.cookies = {}
+        self.default_headers = default_headers
+        self.sessions = {}
+        self.session_regex = session_regex and re.compile(session_regex)
 
     def get(self,url,cookies=None,headers=None,auth=None):
         return self.post(url,data=None,cookies=cookies,headers=headers)
 
     def post(self,url,data=None,cookies=None,headers=None,auth=None):
         self.url = self.app+url
+        
+        # if this POST form requires a postback do it
         if data and '_formname' in data and self.postbacks and \
                 self.history and self.history[-1][1]!=self.url:
             # to bypass the web2py CSRF need to get formkey
             # before submitting the form
-            self.get(url,cookies=None,headers=None,auth=None)
+            self.get(url,cookies=cookies,headers=headers,auth=auth)
-        if cookies is None: cookies = self.cookies
+
+        # unless cookies are specified, recycle cookies
+        if cookies is None:
+            cookies = self.cookies
+        cookies = cookies or {}
+        headers = headers or {}
+
+        # if required do basic auth
         if auth:
             auth_handler = urllib2.HTTPBasicAuthHandler()
             auth_handler.add_password(**auth)
             opener = urllib2.build_opener(auth_handler)
         else:
             opener = urllib2.build_opener()
+
         # copy headers from dict to list of key,value
         headers_list = []
-        for key,value in (headers or {}).iteritems():
+        for key,value in self.default_headers.iteritems():
+            if not key in headers:
+                headers[key] = value
+        for key,value in headers.iteritems():
             if isinstance(value,(list,tuple)):
                 for v in value: headers_list.append((key,v))
             else:
                 headers_list.append((key,value))
+
         # move cookies to headers
-        for key,value in (cookies or {}).iteritems():
+        for key,value in cookies.iteritems():
             headers_list.append(('Cookie','%s=%s' % (key,value)))
+
         # add headers to request
         for key,value in headers_list:
             opener.addheaders.append((key,str(value)))
-        if data is not None:
+
-            self.method = 'POST'
+        # assume everything is ok and make http request
-            # if there is only one form, set _formname automatically
+        error = None
-            if not '_formname' in data and len(self.forms)==1:
+        try:
-                data['_formname'] = self.forms.keys()[0]
+            if data is not None:
-            # if there is no formkey but it is known, set it
+                self.method = 'POST'
-            if '_formname' in data and not '_formkey' in data and \
+
-                    data['_formname'] in self.forms:
+                # if there is only one form, set _formname automatically
-                data['_formkey'] = self.forms[data['_formname']]
+                if not '_formname' in data and len(self.forms)==1:
-            data = urllib.urlencode(data)
+                    data['_formname'] = self.forms.keys()[0]
-            t0 = time.time()
+
-            self.request = opener.open(self.url,data)
+                # if there is no formkey but it is known, set it
+                if '_formname' in data and not '_formkey' in data and \
+                        data['_formname'] in self.forms:
+                    data['_formkey'] = self.forms[data['_formname']]
+
+                # time the POST request
+                data = urllib.urlencode(data)
+                t0 = time.time()
+                self.response = opener.open(self.url,data)
+                self.time = time.time()-t0
+            else:
+                self.method = 'GET'
+
+                # time the GET request
+                t0 = time.time()
+                self.response = opener.open(self.url)
+                self.time = time.time()-t0
+        except urllib2.HTTPError, error:
+            # catch HTTP errors
             self.time = time.time()-t0
-        else:
+            self.response = error
-            self.method = 'GET'
+
-            t0 = time.time()
+        self.status = self.response.getcode()
-            self.request = opener.open(self.url)
+        self.text = self.response.read()
-            self.time = time.time()-t0
+        self.headers = dict(self.response.headers)
-        self.status = self.request.getcode()
+
-        self.text = self.request.read()        
+        # treat web2py tickets as special types of errors
-        self.headers = dict(self.request.headers)
+        if error is not None:
+            if 'web2py_error' in self.headers:
+                raise RuntimeError, self.headers['web2py_error']
+            else:
+                raise error
+
         # parse headers into cookies
-        self.cookies = dict(item[:item.find(';')].split('=') for item in \
+        if 'set-cookie' in self.headers:
-                                self.headers.get('set-cookie','').split(','))
+            self.cookies = dict(
+                item[:item.find(';')].split('=') for item in \
+                    self.headers['set-cookie'].split(','))
+        else:
+            self.cookies = {}
+
+        # check is a new session id has been issued, symptom of broken session
+        if self.session_regex is not None:
+            for cookie, value in self.cookies.iteritems():
+                match = self.session_regex.match(cookie)
+                if match:
+                    name = match.group('name')
+                    if name in self.sessions and self.sessions[name]!=value:
+                        raise RuntimeError, 'Broken sessions %s' % name
+                    self.sessions[name] = value
+
+        # find all forms and formkeys in page
         self.forms = {}
-        # find all forms and formkeys
+        for match in FORM_REGEX.finditer(self.text):
-        for match in WebClient.regex.finditer(self.text):
             self.forms[match.group('formname')] = match.group('formkey')
+
         # log this request
         self.history.append((self.method,self.url,self.status,self.time))
 
 def test_web2py_registration_and_login():
+    # from gluon.contrib.webclient import WebClient
+    # start a web2py instance for testing
+
     client = WebClient('http://127.0.0.1:8000/welcome/default/')
     client.get('index')
-    session_id_welcome = client.cookies['session_id_welcome']
 
+    # register
     data = dict(first_name = 'Homer',
                 last_name = 'Simpson',
                 email = 'homer@web2py.com',
@@ -95,22 +176,29 @@ def test_web2py_registration_and_login():
                 _formname = 'register')
     client.post('user/register',data = data)
 
+    # logout
+    client.get('user/logout')
+
+    # login
     data = dict(email='homer@web2py.com',
                 password='test',
                 _formname = 'login')
     client.post('user/login',data = data)
 
-    client.get('index')
-
     # check registration and login were successful
+    client.get('user/profile')
     assert 'Welcome Homer' in client.text
 
-    # check we are always in the same session
+    # print some variables
-    assert session_id_welcome == client.cookies['session_id_welcome']
+    print '\nsessions:\n',client.sessions
-
+    print '\nheaders:\n',client.headers
+    print '\ncookies:\n',client.cookies
+    print '\nforms:\n',client.forms
+    print
     for method, url, status, t in client.history:
         print method, url, status, t
 
 if __name__ == '__main__':
     test_web2py_registration_and_login()
 
+

gluon/tests/test_cache.py

@@ -70,3 +70,5 @@ if __name__ == '__main__':
     setUpModule()       # pre-python-2.7
     unittest.main()
     tearDownModule()
+
+

gluon/tests/test_contribs.py

@@ -39,3 +39,4 @@ class TestContribs(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
+

gluon/tests/test_dal.py

@@ -579,3 +579,4 @@ class TestImportExportUuidFields(unittest.TestCase):
 if __name__ == '__main__':
     unittest.main()
     tearDownModule()
+

gluon/tests/test_html.py

@@ -209,3 +209,4 @@ class TestBareHelpers(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
+

gluon/tests/test_is_url.py

@@ -642,3 +642,4 @@ class TestUnicode(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
+

gluon/tests/test_languages.py

@@ -92,3 +92,4 @@ except ImportError:
 
 if __name__ == '__main__':
     unittest.main()
+

gluon/tests/test_router.py

@@ -1053,3 +1053,4 @@ if __name__ == '__main__':
     setUpModule()       # pre-2.7
     unittest.main()
     tearDownModule()
+

gluon/tests/test_routes.py

@@ -368,3 +368,4 @@ if __name__ == '__main__':
     setUpModule()       # pre-2.7
     unittest.main()
     tearDownModule()
+

gluon/tests/test_storage.py

@@ -76,3 +76,4 @@ class TestStorage(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
+

gluon/tests/test_template.py

@@ -58,3 +58,4 @@ class TestVirtualFields(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
+

gluon/tests/test_utils.py

@@ -25,3 +25,4 @@ class TestUtils(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
+

gluon/tests/test_web.py

@@ -18,8 +18,8 @@ class TestWeb(unittest.TestCase):
         client = WebClient('http://127.0.0.1:8000/welcome/default/')
 
         client.get('index')
-        session_id_welcome = client.cookies['session_id_welcome']
 
+        # register
         data = dict(first_name = 'Homer',
                     last_name = 'Simpson',
                     email = 'homer@web2py.com',
@@ -28,20 +28,19 @@ class TestWeb(unittest.TestCase):
                     _formname = 'register')
         client.post('user/register',data = data)
 
+        # logout
+        client.get('user/logout')
 
+        # login again
         data = dict(email='homer@web2py.com',
                     password='test',
                     _formname = 'login')
         client.post('user/login',data = data)
 
-        client.get('index')
-
         # check registration and login were successful
+        client.get('index')
         self.assertTrue('Welcome Homer' in client.text)
 
-        # check we are always in the same session
-        self.assertEqual(session_id_welcome,
-                         client.cookies['session_id_welcome'])
-
 if __name__ == '__main__':
     unittest.main()
+