Merge pull request sigmavirus24#321 from zubchick/develop

sigmavirus24 · sigmavirus24 · commit cc8248f23298 · 2014-12-05T21:23:21.000-06:00
issue with pickle in GitHubSession
diff --git a/github3/session.py b/github3/session.py
@@ -19,6 +19,7 @@ def requires_2fa(response):
 
 class GitHubSession(requests.Session):
     auth = None
+    __attrs__ = requests.Session.__attrs__ + ['base_url', 'two_factor_auth_cb']
 
     def __init__(self):
         super(GitHubSession, self).__init__()
diff --git a/tests/unit/test_github_session.py b/tests/unit/test_github_session.py
@@ -1,3 +1,8 @@
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
+
 import pytest
 
 import requests
@@ -231,3 +236,14 @@ def test_retrieve_client_credentials_returns_none(self):
         """
         s = self.build_session()
         assert s.retrieve_client_credentials() == (None, None)
+
+    def test_pickling(self):
+        s = self.build_session('https://api.github.com')
+        dumped = pickle.dumps(s, pickle.HIGHEST_PROTOCOL)
+        loaded = pickle.loads(dumped)
+
+        assert hasattr(loaded, 'base_url')
+        assert hasattr(loaded, 'two_factor_auth_cb')
+
+        assert loaded.base_url == s.base_url
+        assert loaded.two_factor_auth_cb == s.two_factor_auth_cb
