Add GitHub#temporary_basic_auth context manager

sigmavirus24 · sigmavirus24 · commit fe84d7733414 · 2014-02-21T16:55:26.000-06:00
- Included: Tests!
diff --git a/github3/session.py b/github3/session.py
@@ -4,6 +4,7 @@
 from collections import Callable
 from github3 import __version__
 from logging import getLogger
+from contextlib import contextmanager
 
 __url_cache__ = {}
 __logs__ = getLogger(__package__)
@@ -109,3 +110,15 @@ def token_auth(self, token):
             })
         # Unset username/password so we stop sending them
         self.auth = None
+
+    @contextmanager
+    def temporary_basic_auth(self, *auth):
+        old_basic_auth = self.auth
+        old_token_auth = self.headers.get('Authorization')
+
+        self.basic_auth(*auth)
+        yield
+
+        self.auth = old_basic_auth
+        if old_token_auth:
+            self.headers['Authorization'] = old_token_auth
diff --git a/tests/unit/test_github_session.py b/tests/unit/test_github_session.py
@@ -184,3 +184,19 @@ def test_issubclass_of_requests_Session(self):
         """Test that GitHubSession is a subclass of requests.Session"""
         assert issubclass(session.GitHubSession,
                           requests.Session)
+
+    def test_can_use_temporary_basic_auth(self):
+        """Test that temporary_basic_auth resets old auth."""
+        s = self.build_session()
+        s.basic_auth('foo', 'bar')
+        with s.temporary_basic_auth('temp', 'pass'):
+            assert s.auth != ('foo', 'bar')
+
+        assert s.auth == ('foo', 'bar')
+
+    def test_temporary_basic_auth_replaces_auth(self):
+        """Test that temporary_basic_auth sets the proper credentials."""
+        s = self.build_session()
+        s.basic_auth('foo', 'bar')
+        with s.temporary_basic_auth('temp', 'pass'):
+            assert s.auth == ('temp', 'pass')
