Merge pull request #5 from discogs/headers

Add headers keyword to request methods.

Author: josiah-wolf-oberholtzer

MANIFEST.in

@@ -3,7 +3,8 @@ include LICENSE
 include requirements.txt
 include setup.py
 include MANIFEST.in
-include test*
+include tests/*
 include tox.ini
+include pytest.ini
 recursive-include cas_client *
 recursive-exclude cas_client *.pyc

Makefile

@@ -0,0 +1,9 @@
+clean:
+	find . -name '*.pyc' | xargs rm -Rif
+	find . -name '*egg-info' | xargs rm -Rif
+	find . -name __pycache__ | xargs rm -Rif
+	rm -Rif dist/
+	rm -Rif build/
+
+package:
+	python setup.py build sdist

cas_client/_version.py

@@ -1,2 +1,2 @@
-__version_info__ = (0, 1, 3)
+__version_info__ = (0, 1, 4)
 __version__ = '.'.join(str(_) for _ in __version_info__)

cas_client/cas_client.py

@@ -35,6 +35,7 @@ def __init__(
         proxy_callback=None,
         verify_certificates=False,
         session_storage_adapter=None,
+        headers=None,
         ):
         self._auth_prefix = auth_prefix
         self._proxy_callback = proxy_callback
@@ -43,16 +44,17 @@ def __init__(
         self._service_url = service_url
         self._session_storage_adapter = session_storage_adapter
         self._verify_certificates = bool(verify_certificates)
+        self._headers = headers
 
     ### PUBLIC METHODS ###
 
-    def acquire_auth_token_ticket(self):
+    def acquire_auth_token_ticket(self, headers=None):
         '''
         Acquire an auth token from the CAS server.
         '''
         logging.debug('[CAS] Acquiring Auth token ticket')
         url = self._get_auth_token_tickets_url()
-        text = self._perform_post(url)
+        text = self._perform_post(url, headers=headers)
         auth_token_ticket = json.loads(text)['ticket']
         logging.debug('[CAS] Acquire Auth token ticket: {}'.format(
             auth_token_ticket))
@@ -257,6 +259,7 @@ def perform_api_request(
         private_key,
         method='POST',
         service_url=None,
+        headers=None,
         **kwargs
         ):
         '''
@@ -272,34 +275,47 @@ def perform_api_request(
             **kwargs
             )
         if method == 'GET':
-            response = self._perform_get(url)
+            response = self._perform_get(url, headers=headers)
         elif method == 'POST':
-            response = self._perform_post(url)
+            response = self._perform_post(url, headers=headers)
         return response
 
-    def perform_proxy(self, proxy_ticket):
+    def perform_proxy(self, proxy_ticket, headers=None):
         '''
         Fetch a response from the remote CAS `proxy` endpoint.
         '''
         url = self._get_proxy_url(ticket=proxy_ticket)
         logging.debug('[CAS] Proxy URL: {}'.format(url))
-        return self._perform_cas_call(url, ticket=proxy_ticket)
+        return self._perform_cas_call(
+            url,
+            ticket=proxy_ticket,
+            headers=headers,
+            )
 
-    def perform_proxy_validate(self, proxied_service_ticket):
+    def perform_proxy_validate(self, proxied_service_ticket, headers=None):
         '''
         Fetch a response from the remote CAS `proxyValidate` endpoint.
         '''
         url = self._get_proxy_validate_url(ticket=proxied_service_ticket)
         logging.debug('[CAS] ProxyValidate URL: {}'.format(url))
-        return self._perform_cas_call(url, ticket=proxied_service_ticket)
+        return self._perform_cas_call(
+            url,
+            ticket=proxied_service_ticket,
+            headers=headers,
+            )
 
-    def perform_service_validate(self, ticket=None, service_url=None):
+    def perform_service_validate(
+        self,
+        ticket=None,
+        service_url=None,
+        headers=None,
+        ):
         '''
         Fetch a response from the remote CAS `serviceValidate` endpoint.
         '''
         url = self._get_service_validate_url(ticket, service_url=service_url)
         logging.debug('[CAS] ServiceValidate URL: {}'.format(url))
-        return self._perform_cas_call(url, ticket=ticket)
+        return self._perform_cas_call(url, ticket=ticket, headers=headers)
 
     def session_exists(self, ticket):
         '''
@@ -405,32 +421,42 @@ def _get_service_validate_url(self, ticket, service_url=None):
             url = '{url}&pgtUrl={proxy_url}'.format(url, self.proxy_url)
         return url
 
-    def _perform_cas_call(self, url, ticket):
+    def _perform_cas_call(self, url, ticket, headers=None):
         if ticket is not None:
             logging.debug('[CAS] Requesting Ticket Validation')
-            response_text = self._perform_get(url)
+            response_text = self._perform_get(url, headers=headers)
             response_text = self._clean_up_response_text(response_text)
             if response_text:
                 logging.debug('[CAS] Response:\n{}'.format(response_text))
                 return CASResponse(response_text)
         logging.debug('[CAS] Response: None')
         return None
 
-    def _perform_get(self, url):
+    def _perform_get(self, url, headers=None):
+        headers = headers or self.headers
         try:
-            response = requests.get(url, verify=self.verify_certificates)
+            response = requests.get(
+                url,
+                verify=self.verify_certificates,
+                headers=headers,
+                )
             return response.text
         except requests.HTTPError:
             return None
 
-    def _perform_post(self, url):
+    def _perform_post(self, url, headers=None):
+        headers = headers or self.headers
         try:
-            response = requests.post(url, verify=self.verify_certificates)
+            response = requests.post(
+                url,
+                verify=self.verify_certificates,
+                headers=headers,
+                )
             return response.text
         except requests.HTTPError:
             return None
 
-    ### PUBLIC METHODS ###
+    ### PUBLIC PROPERTIES ###
 
     @property
     def auth_prefix(self):
@@ -439,6 +465,10 @@ def auth_prefix(self):
         '''
         return self._auth_prefix
 
+    @property
+    def headers(self):
+        return self._headers
+
     @property
     def proxy_callback(self):
         '''

pytest.ini

@@ -0,0 +1,9 @@
+[pytest]
+addopts =
+    --doctest-modules
+    -vv
+doctest_optionflags =
+    ELLIPSIS
+    NORMALIZE_WHITESPACE
+python_files = test*.py
+testpaths = tests

test.py renamed to tests/test.py

@@ -1,10 +1,15 @@
 # -*- encoding: utf-8 -*-
+import os
 import unittest
 from cas_client import CASClient, CASResponse
 try:
     from urlparse import parse_qs
 except ImportError:
     from urllib.parse import parse_qs
+try:
+    import mock
+except ImportError:
+    from unittest import mock
 
 
 class TestCase(unittest.TestCase):
@@ -49,6 +54,11 @@ class TestCase(unittest.TestCase):
     </samlp:LogoutRequest>
     """
 
+    private_key_filepath = os.path.join(
+        os.path.abspath(os.path.dirname(__file__)),
+        'test_private_key.pem',
+        )
+
     def test_success(self):
         response = CASResponse(self.response_text)
         self.assertTrue(response.success)
@@ -70,11 +80,13 @@ def test_failure(self):
 
     def test_perform_service_validate(self):
         cas_client = CASClient('https://dummy.url')
-        cas_client._perform_get = lambda url: self.response_text
-        response = cas_client.perform_service_validate(
-            ticket='FOO',
-            service_url='BAR',
-            )
+        assert not cas_client.headers
+        with mock.patch('cas_client.CASClient._perform_get') as m:
+            m.return_value = self.response_text
+            response = cas_client.perform_service_validate(
+                ticket='FOO',
+                service_url='BAR',
+                )
         self.assertTrue(response.success)
         self.assertEqual(response.attributes, {
             u'i2a2characteristics': u'0,3592,2000',
@@ -87,6 +99,43 @@ def test_perform_service_validate(self):
         self.assertEqual(response.response_type, 'authenticationSuccess')
         self.assertEqual(response.user, 'jott')
 
+    def test_perform_service_validate_headers_call(self):
+        class MockResponse(object):
+            text = self.response_text
+
+        cas_client = CASClient('https://dummy.url')
+        assert not cas_client.headers
+        with mock.patch('requests.get') as m:
+            m.return_value = MockResponse()
+            cas_client.perform_service_validate(
+                ticket='FOO',
+                service_url='BAR',
+                headers={'baz': 'quux'},
+                )
+        m.assert_called_with(
+            'https://dummy.url/cas/serviceValidate?ticket=FOO&service=BAR',
+            headers={'baz': 'quux'},
+            verify=False,
+            )
+
+    def test_perform_service_validate_headers_init(self):
+        class MockResponse(object):
+            text = self.response_text
+
+        cas_client = CASClient('https://dummy.url', headers={'baz': 'quux'})
+        assert cas_client.headers == {'baz': 'quux'}
+        with mock.patch('requests.get') as m:
+            m.return_value = MockResponse()
+            cas_client.perform_service_validate(
+                ticket='FOO',
+                service_url='BAR',
+                )
+        m.assert_called_with(
+            'https://dummy.url/cas/serviceValidate?ticket=FOO&service=BAR',
+            headers={'baz': 'quux'},
+            verify=False,
+            )
+
     def test_get_destroy_other_sessions_url(self):
         cas_client = CASClient('https://dummy.url')
         service_url = 'https://app.url'
@@ -134,8 +183,7 @@ def test_get_api_url(self):
         api_resource = 'do_something_useful'
         auth_token_ticket = 'ATT-1234'
         authenticator = 'my_company_ldap'
-        private_key_filepath = 'test_private_key.pem'
-        with open(private_key_filepath, 'r') as file_pointer:
+        with open(self.private_key_filepath, 'r') as file_pointer:
             private_key = file_pointer.read()
         service_url = 'https://example.com'
         kwargs = {
@@ -184,8 +232,7 @@ def test_get_auth_token_login_url(self):
         authenticator = 'my_company_ldap'
         username = 'my_user'
         service_url = 'https://example.com'
-        private_key_filepath = 'test_private_key.pem'
-        with open(private_key_filepath, 'r') as file_pointer:
+        with open(self.private_key_filepath, 'r') as file_pointer:
             private_key = file_pointer.read()
         url = cas_client.get_auth_token_login_url(
             auth_token_ticket=auth_token_ticket,
@@ -220,3 +267,48 @@ def test_get_auth_token_login_url(self):
                 ),
             'service': 'https://example.com',
             }
+
+    def test_acquire_auth_token_ticket_no_headers(self):
+        class MockResponse(object):
+            text = '{"ticket": "FOO"}'
+
+        cas_client = CASClient('https://dummy.url')
+        assert not cas_client.headers
+        with mock.patch('requests.post') as m:
+            m.return_value = MockResponse()
+            cas_client.acquire_auth_token_ticket()
+        m.assert_called_with(
+            'https://dummy.url/cas/api/auth_token_tickets',
+            headers=None,
+            verify=False,
+            )
+
+    def test_acquire_auth_token_ticket_headers_call(self):
+        class MockResponse(object):
+            text = '{"ticket": "FOO"}'
+
+        cas_client = CASClient('https://dummy.url')
+        assert not cas_client.headers
+        with mock.patch('requests.post') as m:
+            m.return_value = MockResponse()
+            cas_client.acquire_auth_token_ticket(headers={'baz': 'quux'})
+        m.assert_called_with(
+            'https://dummy.url/cas/api/auth_token_tickets',
+            headers={'baz': 'quux'},
+            verify=False,
+            )
+
+    def test_acquire_auth_token_ticket_headers_init(self):
+        class MockResponse(object):
+            text = '{"ticket": "FOO"}'
+
+        cas_client = CASClient('https://dummy.url', headers={'baz': 'quux'})
+        assert cas_client.headers == {'baz': 'quux'}
+        with mock.patch('requests.post') as m:
+            m.return_value = MockResponse()
+            cas_client.acquire_auth_token_ticket()
+        m.assert_called_with(
+            'https://dummy.url/cas/api/auth_token_tickets',
+            headers={'baz': 'quux'},
+            verify=False,
+            )

test_private_key.pem renamed to tests/test_private_key.pem

@@ -1,8 +1,13 @@
 [tox]
-envlist = py27, py34
+envlist=
+    py27
+    py34
+    py35
+    py36
 
 [testenv]
-deps=pytest
+deps=
+    pytest
+    py27: mock
 commands=
-    py.test -rf -vv test.py
-    python -m doctest cas_client/cas_client.py
+    py.test