Switch django models from pickle to jsonpickle (googleapis#614)

Author: waprin

oauth2client/contrib/django_util/views.py

@@ -22,13 +22,13 @@
 import hashlib
 import json
 import os
-import pickle
 
 from django import http
 from django import shortcuts
 from django.conf import settings
 from django.core import urlresolvers
 from django.shortcuts import redirect
+import jsonpickle
 from six.moves.urllib import parse
 
 from oauth2client import client
@@ -71,7 +71,7 @@ def _make_flow(request, scopes, return_url=None):
             urlresolvers.reverse("google_oauth:callback")))
 
     flow_key = _FLOW_KEY.format(csrf_token)
-    request.session[flow_key] = pickle.dumps(flow)
+    request.session[flow_key] = jsonpickle.encode(flow)
     return flow
 
 
@@ -89,7 +89,7 @@ def _get_flow_for_token(csrf_token, request):
         CSRF token.
     """
     flow_pickle = request.session.get(_FLOW_KEY.format(csrf_token), None)
-    return None if flow_pickle is None else pickle.loads(flow_pickle)
+    return None if flow_pickle is None else jsonpickle.decode(flow_pickle)
 
 
 def oauth2_callback(request):

tests/contrib/django_util/test_views.py

@@ -152,8 +152,8 @@ def setUp(self):
         self.user = django_models.User.objects.create_user(
             username='bill', email='[email protected]', password='hunter2')
 
-    @mock.patch('oauth2client.contrib.django_util.views.pickle')
-    def test_callback_works(self, pickle):
+    @mock.patch('oauth2client.contrib.django_util.views.jsonpickle')
+    def test_callback_works(self, jsonpickle_mock):
         request = self.factory.get('oauth2/oauth2callback', data={
             'state': json.dumps(self.fake_state),
             'code': 123
@@ -169,9 +169,10 @@ def test_callback_works(self, pickle):
             redirect_uri=request.build_absolute_uri("oauth2/oauth2callback"))
 
         name = 'google_oauth2_flow_{0}'.format(self.CSRF_TOKEN)
-        self.session[name] = pickle.dumps(flow)
+        pickled_flow = object()
+        self.session[name] = pickled_flow
         flow.step2_exchange = mock.Mock()
-        pickle.loads.return_value = flow
+        jsonpickle_mock.decode.return_value = flow
 
         request.session = self.session
         request.user = self.user
@@ -180,9 +181,10 @@ def test_callback_works(self, pickle):
         self.assertEqual(
             response.status_code, django.http.HttpResponseRedirect.status_code)
         self.assertEqual(response['Location'], self.RETURN_URL)
+        jsonpickle_mock.decode.assert_called_once_with(pickled_flow)
 
-    @mock.patch('oauth2client.contrib.django_util.views.pickle')
-    def test_callback_handles_bad_flow_exchange(self, pickle):
+    @mock.patch('oauth2client.contrib.django_util.views.jsonpickle')
+    def test_callback_handles_bad_flow_exchange(self, jsonpickle_mock):
         request = self.factory.get('oauth2/oauth2callback', data={
             "state": json.dumps(self.fake_state),
             "code": 123
@@ -198,17 +200,19 @@ def test_callback_handles_bad_flow_exchange(self, pickle):
             redirect_uri=request.build_absolute_uri('oauth2/oauth2callback'))
 
         session_key = 'google_oauth2_flow_{0}'.format(self.CSRF_TOKEN)
-        self.session[session_key] = pickle.dumps(flow)
+        pickled_flow = object()
+        self.session[session_key] = pickled_flow
 
         def local_throws(code):
             raise client.FlowExchangeError('test')
 
         flow.step2_exchange = local_throws
-        pickle.loads.return_value = flow
+        jsonpickle_mock.decode.return_value = flow
 
         request.session = self.session
         response = views.oauth2_callback(request)
         self.assertIsInstance(response, http.HttpResponseBadRequest)
+        jsonpickle_mock.decode.assert_called_once_with(pickled_flow)
 
     def test_error_returns_bad_request(self):
         request = self.factory.get('oauth2/oauth2callback', data={

tox.ini

@@ -14,6 +14,7 @@ basedeps = mock>=1.3.0
 deps = {[testenv]basedeps}
        django
        keyring
+       jsonpickle
 setenv =
     pypy: with_gmp=no
     DJANGO_SETTINGS_MODULE=tests.contrib.django_util.settings