Added a test to improve coverage and cleaned up some existing tests.

3 files changed, 30 insertions, 11 deletions

diff --git a/jwt/api.py b/jwt/api.py
index fd0dc0f..e21ed00 100644
--- a/jwt/api.py
+++ b/jwt/api.py
@@ -84,7 +84,7 @@ class PyJWT(object):
 
         # Header
         header = {'typ': 'JWT', 'alg': algorithm}
-        
+
         if headers:
             header.update(headers)
 
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index 8eb0eb3..5eb24bc 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -113,9 +113,9 @@ class TestAlgorithms(unittest.TestCase):
     def test_rsa_verify_should_return_false_if_signature_invalid(self):
         algo = RSAAlgorithm(RSAAlgorithm.SHA256)
 
-        jwt_message = ensure_bytes('Hello World!')
+        message = ensure_bytes('Hello World!')
 
-        jwt_sig = base64.b64decode(ensure_bytes(
+        sig = base64.b64decode(ensure_bytes(
             'yS6zk9DBkuGTtcBzLUzSpo9gGJxJFOGvUqN01iLhWHrzBQ9ZEz3+Ae38AXp'
             '10RWwscp42ySC85Z6zoN67yGkLNWnfmCZSEv+xqELGEvBJvciOKsrhiObUl'
             '2mveSc1oeO/2ujkGDkkkJ2epn0YliacVjZF5+/uDmImUfAAj8lzjnHlzYix'
@@ -123,21 +123,21 @@ class TestAlgorithms(unittest.TestCase):
             'fHJnNUzAEUOXS0WahHVb57D30pcgIji9z923q90p5c7E2cU8V+E1qe8NdCA'
             'APCDzZZ9zQ/dgcMVaBrGrgimrcLbPjueOKFgSO+SSjIElKA=='))
 
-        jwt_sig += ensure_bytes('123')  # Signature is now invalid
+        sig += ensure_bytes('123')  # Signature is now invalid
 
         with open(key_path('testkey_rsa.pub'), 'r') as keyfile:
-            jwt_pub_key = algo.prepare_key(keyfile.read())
+            pub_key = algo.prepare_key(keyfile.read())
 
-        result = algo.verify(jwt_message, jwt_pub_key, jwt_sig)
+        result = algo.verify(message, pub_key, sig)
         self.assertFalse(result)
 
     @unittest.skipIf(not has_crypto, 'Not supported without cryptography library')
     def test_rsa_verify_should_return_true_if_signature_valid(self):
         algo = RSAAlgorithm(RSAAlgorithm.SHA256)
 
-        jwt_message = ensure_bytes('Hello World!')
+        message = ensure_bytes('Hello World!')
 
-        jwt_sig = base64.b64decode(ensure_bytes(
+        sig = base64.b64decode(ensure_bytes(
             'yS6zk9DBkuGTtcBzLUzSpo9gGJxJFOGvUqN01iLhWHrzBQ9ZEz3+Ae38AXp'
             '10RWwscp42ySC85Z6zoN67yGkLNWnfmCZSEv+xqELGEvBJvciOKsrhiObUl'
             '2mveSc1oeO/2ujkGDkkkJ2epn0YliacVjZF5+/uDmImUfAAj8lzjnHlzYix'
@@ -146,9 +146,9 @@ class TestAlgorithms(unittest.TestCase):
             'APCDzZZ9zQ/dgcMVaBrGrgimrcLbPjueOKFgSO+SSjIElKA=='))
 
         with open(key_path('testkey_rsa.pub'), 'r') as keyfile:
-            jwt_pub_key = algo.prepare_key(keyfile.read())
+            pub_key = algo.prepare_key(keyfile.read())
 
-        result = algo.verify(jwt_message, jwt_pub_key, jwt_sig)
+        result = algo.verify(message, pub_key, sig)
         self.assertTrue(result)
 
     @unittest.skipIf(not has_crypto, 'Not supported without cryptography library')
diff --git a/tests/test_api.py b/tests/test_api.py
index 0afb8bd..4d5f415 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -13,8 +13,9 @@ from jwt.exceptions import (
     InvalidAlgorithmError, InvalidAudienceError, InvalidIssuedAtError,
     InvalidIssuerError
 )
+from jwt.utils import base64url_decode
 
-from .compat import text_type, unittest
+from .compat import string_types, text_type, unittest
 from .utils import ensure_bytes
 
 try:
@@ -856,6 +857,24 @@ class TestAPI(unittest.TestCase):
         payload = self.jwt.decode(token, 'secret')
         self.assertEqual(payload, {'some_decimal': 'it worked'})
 
+    def test_encode_headers_parameter_adds_headers(self):
+        headers = {'testheader': True}
+        token = self.jwt.encode({'msg': 'hello world'}, 'secret', headers=headers)
+
+        if not isinstance(token, string_types):
+            token = token.decode()
+
+        header = token[0:token.index('.')].encode()
+        header = base64url_decode(header)
+
+        if not isinstance(header, text_type):
+            header = header.decode()
+
+        header_obj = json.loads(header)
+
+        self.assertIn('testheader', header_obj)
+        self.assertEqual(header_obj['testheader'], headers['testheader'])
+
 
 if __name__ == '__main__':
     unittest.main()