From 4ca55219d175a78ee120cbd271c816032811c1e5 Mon Sep 17 00:00:00 2001
From: Yasuhisa Yamazaki <1@yasu21.net>
Date: Wed, 18 Jun 2014 17:38:52 +0900
Subject: [PATCH] Fix #41

---
 django_mobile/cache/__init__.py   | 17 ++++++++++-------
 django_mobile/cache/middleware.py |  8 +++++++-
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/django_mobile/cache/__init__.py b/django_mobile/cache/__init__.py
index 9ae9229..2356fbf 100644
--- a/django_mobile/cache/__init__.py
+++ b/django_mobile/cache/__init__.py
@@ -1,23 +1,26 @@
-from functools import wraps
 from django.views.decorators.cache import cache_page as _cache_page
 from django.utils.decorators import decorator_from_middleware
-from django_mobile.cache.middleware import CacheFlavourMiddleware
 
+from django_mobile.cache.middleware import CacheFlavourMiddleware, CacheFlavourRequestMiddleware, \
+    CacheFlavourResponseMiddleware
 
-__all__ = ('cache_page', 'vary_on_flavour')
 
+__all__ = ('cache_page', 'vary_on_flavour')
 
 vary_on_flavour = decorator_from_middleware(CacheFlavourMiddleware)
+vary_on_flavour_request = decorator_from_middleware(CacheFlavourRequestMiddleware)
+vary_on_flavour_response = decorator_from_middleware(CacheFlavourResponseMiddleware)
 
 
 def cache_page(*args, **kwargs):
     '''
     Same as django's ``cache_page`` decorator, but wraps the view into
-    ``vary_on_flavour`` decorator before. Makes it possible to serve multiple
-    flavours without getting into trouble with django's caching that doesn't
+    ``vary_on_flavour_request`` and ``vary_on_flavour_response`` decorators.
+    Makes it possible to serve multiple flavours without getting into trouble with django's caching that doesn't
     know about flavours.
     '''
-    decorator = _cache_page(*args, **kwargs)
+
     def flavoured_decorator(func):
-        return decorator(vary_on_flavour(func))
+        return vary_on_flavour_request(_cache_page(*args, **kwargs)(vary_on_flavour_response(func)))
+
     return flavoured_decorator
diff --git a/django_mobile/cache/middleware.py b/django_mobile/cache/middleware.py
index 1bf9f5d..4736aa8 100644
--- a/django_mobile/cache/middleware.py
+++ b/django_mobile/cache/middleware.py
@@ -2,10 +2,16 @@
 from django.utils.cache import patch_vary_headers
 
 
-class CacheFlavourMiddleware(object):
+class CacheFlavourRequestMiddleware(object):
     def process_request(self, request):
         _set_request_header(request, get_flavour(request))
 
+
+class CacheFlavourResponseMiddleware(object):
     def process_response(self, request, response):
         patch_vary_headers(response, ['X-Flavour'])
         return response
+
+
+class CacheFlavourMiddleware(CacheFlavourRequestMiddleware, CacheFlavourResponseMiddleware):
+    pass
\ No newline at end of file