1 import urllib

     2 import sys

     3 from cStringIO import StringIO

     4 from django.conf import settings

     5 from django.contrib.auth import authenticate, login

     6 from django.core.handlers.base import BaseHandler

     7 from django.core.handlers.wsgi import WSGIRequest

     8 from django.core.signals import got_request_exception

     9 from django.dispatch import dispatcher

    10 from django.http import SimpleCookie, HttpRequest

    11 from django.template import TemplateDoesNotExist

    12 from django.test import signals

    13 from django.utils.functional import curry

    14 from django.utils.encoding import smart_str

    15 from django.utils.http import urlencode

    16 from django.utils.itercompat import is_iterable

    17 

    18 BOUNDARY = 'BoUnDaRyStRiNg'

    19 MULTIPART_CONTENT = 'multipart/form-data; boundary=%s' % BOUNDARY

    20 

    21 class ClientHandler(BaseHandler):

    22     """

    23     A HTTP Handler that can be used for testing purposes.

    24     Uses the WSGI interface to compose requests, but returns

    25     the raw HttpResponse object

    26     """

    27     def __call__(self, environ):

    28         from django.conf import settings

    29         from django.core import signals

    30 

    31         # Set up middleware if needed. We couldn't do this earlier, because

    32         # settings weren't available.

    33         if self._request_middleware is None:

    34             self.load_middleware()

    35 

    36         dispatcher.send(signal=signals.request_started)

    37         try:

    38             request = WSGIRequest(environ)

    39             response = self.get_response(request)

    40 

    41             # Apply response middleware

    42             for middleware_method in self._response_middleware:

    43                 response = middleware_method(request, response)

    44             response = self.apply_response_fixes(request, response)

    45         finally:

    46             dispatcher.send(signal=signals.request_finished)

    47 

    48         return response

    49 

    50 def store_rendered_templates(store, signal, sender, template, context):

    51     "A utility function for storing templates and contexts that are rendered"

    52     store.setdefault('template',[]).append(template)

    53     store.setdefault('context',[]).append(context)

    54 

    55 def encode_multipart(boundary, data):

    56     """

    57     A simple method for encoding multipart POST data from a dictionary of

    58     form values.

    59 

    60     The key will be used as the form data name; the value will be transmitted

    61     as content. If the value is a file, the contents of the file will be sent

    62     as an application/octet-stream; otherwise, str(value) will be sent.

    63     """

    64     lines = []

    65     to_str = lambda s: smart_str(s, settings.DEFAULT_CHARSET)

    66     for (key, value) in data.items():

    67         if isinstance(value, file):

    68             lines.extend([

    69                 '--' + boundary,

    70                 'Content-Disposition: form-data; name="%s"; filename="%s"' % (to_str(key), to_str(value.name)),

    71                 'Content-Type: application/octet-stream',

    72                 '',

    73                 value.read()

    74             ])

    75         else:

    76             if not isinstance(value, basestring) and is_iterable(value):

    77                 for item in value:

    78                     lines.extend([

    79                         '--' + boundary,

    80                         'Content-Disposition: form-data; name="%s"' % to_str(key),

    81                         '',

    82                         to_str(item)

    83                     ])

    84             else:

    85                 lines.extend([

    86                     '--' + boundary,

    87                     'Content-Disposition: form-data; name="%s"' % to_str(key),

    88                     '',

    89                     to_str(value)

    90                 ])

    91 

    92     lines.extend([

    93         '--' + boundary + '--',

    94         '',

    95     ])

    96     return '\r\n'.join(lines)

    97 

    98 class Client:

    99     """

   100     A class that can act as a client for testing purposes.

   101 

   102     It allows the user to compose GET and POST requests, and

   103     obtain the response that the server gave to those requests.

   104     The server Response objects are annotated with the details

   105     of the contexts and templates that were rendered during the

   106     process of serving the request.

   107 

   108     Client objects are stateful - they will retain cookie (and

   109     thus session) details for the lifetime of the Client instance.

   110 

   111     This is not intended as a replacement for Twill/Selenium or

   112     the like - it is here to allow testing against the

   113     contexts and templates produced by a view, rather than the

   114     HTML rendered to the end-user.

   115     """

   116     def __init__(self, **defaults):

   117         self.handler = ClientHandler()

   118         self.defaults = defaults

   119         self.cookies = SimpleCookie()

   120         self.exc_info = None

   121 

   122     def store_exc_info(self, *args, **kwargs):

   123         """

   124         Utility method that can be used to store exceptions when they are

   125         generated by a view.

   126         """

   127         self.exc_info = sys.exc_info()

   128 

   129     def _session(self):

   130         "Obtain the current session variables"

   131         if 'django.contrib.sessions' in settings.INSTALLED_APPS:

   132             engine = __import__(settings.SESSION_ENGINE, {}, {}, [''])

   133             cookie = self.cookies.get(settings.SESSION_COOKIE_NAME, None)

   134             if cookie:

   135                 return engine.SessionStore(cookie.value)

   136         return {}

   137     session = property(_session)

   138 

   139     def request(self, **request):

   140         """

   141         The master request method. Composes the environment dictionary

   142         and passes to the handler, returning the result of the handler.

   143         Assumes defaults for the query environment, which can be overridden

   144         using the arguments to the request.

   145         """

   146 

   147         environ = {

   148             'HTTP_COOKIE':      self.cookies,

   149             'PATH_INFO':         '/',

   150             'QUERY_STRING':      '',

   151             'REQUEST_METHOD':    'GET',

   152             'SCRIPT_NAME':       None,

   153             'SERVER_NAME':       'testserver',

   154             'SERVER_PORT':       80,

   155             'SERVER_PROTOCOL':   'HTTP/1.1',

   156         }

   157         environ.update(self.defaults)

   158         environ.update(request)

   159 

   160         # Curry a data dictionary into an instance of

   161         # the template renderer callback function

   162         data = {}

   163         on_template_render = curry(store_rendered_templates, data)

   164         dispatcher.connect(on_template_render, signal=signals.template_rendered)

   165 

   166         # Capture exceptions created by the handler

   167         dispatcher.connect(self.store_exc_info, signal=got_request_exception)

   168 

   169         try:

   170             response = self.handler(environ)

   171         except TemplateDoesNotExist, e:

   172             # If the view raises an exception, Django will attempt to show

   173             # the 500.html template. If that template is not available,

   174             # we should ignore the error in favor of re-raising the

   175             # underlying exception that caused the 500 error. Any other

   176             # template found to be missing during view error handling

   177             # should be reported as-is.

   178             if e.args != ('500.html',):

   179                 raise

   180 

   181         # Look for a signalled exception and reraise it

   182         if self.exc_info:

   183             raise self.exc_info[1], None, self.exc_info[2]

   184 

   185         # Save the client and request that stimulated the response

   186         response.client = self

   187         response.request = request

   188 

   189         # Add any rendered template detail to the response

   190         # If there was only one template rendered (the most likely case),

   191         # flatten the list to a single element

   192         for detail in ('template', 'context'):

   193             if data.get(detail):

   194                 if len(data[detail]) == 1:

   195                     setattr(response, detail, data[detail][0]);

   196                 else:

   197                     setattr(response, detail, data[detail])

   198             else:

   199                 setattr(response, detail, None)

   200 

   201         # Update persistent cookie data

   202         if response.cookies:

   203             self.cookies.update(response.cookies)

   204 

   205         return response

   206 

   207     def get(self, path, data={}, **extra):

   208         "Request a response from the server using GET."

   209         r = {

   210             'CONTENT_LENGTH':  None,

   211             'CONTENT_TYPE':    'text/html; charset=utf-8',

   212             'PATH_INFO':       urllib.unquote(path),

   213             'QUERY_STRING':    urlencode(data, doseq=True),

   214             'REQUEST_METHOD': 'GET',

   215         }

   216         r.update(extra)

   217 

   218         return self.request(**r)

   219 

   220     def post(self, path, data={}, content_type=MULTIPART_CONTENT, **extra):

   221         "Request a response from the server using POST."

   222 

   223         if content_type is MULTIPART_CONTENT:

   224             post_data = encode_multipart(BOUNDARY, data)

   225         else:

   226             post_data = data

   227 

   228         r = {

   229             'CONTENT_LENGTH': len(post_data),

   230             'CONTENT_TYPE':   content_type,

   231             'PATH_INFO':      urllib.unquote(path),

   232             'REQUEST_METHOD': 'POST',

   233             'wsgi.input':     StringIO(post_data),

   234         }

   235         r.update(extra)

   236 

   237         return self.request(**r)

   238 

   239     def login(self, **credentials):

   240         """Set the Client to appear as if it has sucessfully logged into a site.

   241 

   242         Returns True if login is possible; False if the provided credentials

   243         are incorrect, or the user is inactive, or if the sessions framework is

   244         not available.

   245         """

   246         user = authenticate(**credentials)

   247         if user and user.is_active and 'django.contrib.sessions' in settings.INSTALLED_APPS:

   248             engine = __import__(settings.SESSION_ENGINE, {}, {}, [''])

   249 

   250             # Create a fake request to store login details

   251             request = HttpRequest()

   252             request.session = engine.SessionStore()

   253             login(request, user)

   254 

   255             # Set the cookie to represent the session

   256             self.cookies[settings.SESSION_COOKIE_NAME] = request.session.session_key

   257             self.cookies[settings.SESSION_COOKIE_NAME]['max-age'] = None

   258             self.cookies[settings.SESSION_COOKIE_NAME]['path'] = '/'

   259             self.cookies[settings.SESSION_COOKIE_NAME]['domain'] = settings.SESSION_COOKIE_DOMAIN

   260             self.cookies[settings.SESSION_COOKIE_NAME]['secure'] = settings.SESSION_COOKIE_SECURE or None

   261             self.cookies[settings.SESSION_COOKIE_NAME]['expires'] = None

   262 

   263             # Save the session values

   264             request.session.save()

   265 

   266             return True

   267         else:

   268             return False

   269 

   270     def logout(self):

   271         """Removes the authenticated user's cookies.

   272 

   273         Causes the authenticated user to be logged out.

   274         """

   275         session = __import__(settings.SESSION_ENGINE, {}, {}, ['']).SessionStore()

   276         session.delete(session_key=self.cookies[settings.SESSION_COOKIE_NAME].value)

   277         self.cookies = SimpleCookie()