From 0e134dc9eb7bc09ec31a95d140e9e52818f4985f Mon Sep 17 00:00:00 2001
From: Pablo Klijnjan <pablokl@il.ibm.com>
Date: Sun, 15 Dec 2013 10:25:42 +0200
Subject: [PATCH] Adding utility method to APIClient

    - temporary_credentials acts as the credentials method
      but only for a context
---
 docs/api-guide/testing.md            | 10 ++++++++++
 rest_framework/test.py               | 10 ++++++++++
 rest_framework/tests/test_testing.py | 13 +++++++++++++
 3 files changed, 33 insertions(+)

diff --git a/docs/api-guide/testing.md b/docs/api-guide/testing.md
index 4a8a91682..bfe65253d 100644
--- a/docs/api-guide/testing.md
+++ b/docs/api-guide/testing.md
@@ -158,6 +158,16 @@ Note that calling `credentials` a second time overwrites any existing credential
 
 The `credentials` method is appropriate for testing APIs that require authentication headers, such as basic authentication, OAuth1a and OAuth2 authentication, and simple token authentication schemes.
 
+#### .temporary_credentials(**kwargs)
+
+The `temporary_credentials` method has the same behavior as the `credentials` method but as a context manager. It will
+set the credential for all the requests within the context and then set back the original ones.
+
+    client = APIClient()
+    client.credentials(HTTP_AUTHORIZATION='Token ' + token.key)
+    with client.temporary_credentials(HTTP_AUTHORIZATION='Token ' + other_token.key)
+        # test code within context using this temporary credentials
+
 #### .force_authenticate(user=None, token=None)
 
 Sometimes you may want to bypass authentication, and simple force all requests by the test client to be automatically treated as authenticated.
diff --git a/rest_framework/test.py b/rest_framework/test.py
index 234d10a4a..a84619671 100644
--- a/rest_framework/test.py
+++ b/rest_framework/test.py
@@ -2,7 +2,10 @@
 
 # Note that we import as `DjangoRequestFactory` and `DjangoClient` in order
 # to make it harder for the user to import the wrong thing without realizing.
+
 from __future__ import unicode_literals
+
+from contextlib import contextmanager
 import django
 from django.conf import settings
 from django.test.client import Client as DjangoClient
@@ -127,6 +130,13 @@ class APIClient(APIRequestFactory, DjangoClient):
         """
         self._credentials = kwargs
 
+    @contextmanager
+    def temporary_credentials(self, **kwargs):
+        original_credentials = self._credentials
+        self.credentials(**kwargs)
+        yield
+        self._credentials = original_credentials
+
     def force_authenticate(self, user=None, token=None):
         """
         Forcibly authenticates outgoing requests with the given
diff --git a/rest_framework/tests/test_testing.py b/rest_framework/tests/test_testing.py
index 48b8956b5..27fe8450f 100644
--- a/rest_framework/tests/test_testing.py
+++ b/rest_framework/tests/test_testing.py
@@ -47,6 +47,19 @@ class TestAPITestClient(TestCase):
             response = self.client.get('/view/')
             self.assertEqual(response.data['auth'], 'example')
 
+    def test_temporary_credentials(self):
+        def assert_auth(expected):
+            response = self.client.get('/view/')
+            self.assertEqual(expected, response.data['auth'])
+
+        persistent_credentials = {'HTTP_AUTHORIZATION': 'persistent'}
+        temporal_credentials = {'HTTP_AUTHORIZATION': 'temporal'}
+        self.client.credentials(**persistent_credentials)
+        assert_auth('persistent')
+        with self.client.temporary_credentials(**temporal_credentials):
+            assert_auth('temporal')
+        assert_auth('persistent')
+
     def test_force_authenticate(self):
         """
         Setting `.force_authenticate()` forcibly authenticates each request.