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

Author: mark-adams

jwt/api.py

@@ -84,7 +84,7 @@ def encode(self, payload, key, algorithm='HS256', headers=None, json_encoder=Non
 
         # Header
         header = {'typ': 'JWT', 'alg': algorithm}
-        
+
         if headers:
             header.update(headers)
 

tests/test_algorithms.py

@@ -113,31 +113,31 @@ def test_rsa_should_reject_non_string_key(self):
     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'
             'sn5jGz1H07jYYbi9diixN8IUhXeTafwFg02IcONhum29V40Wu6O5tAKWlJX'
             '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 @@ def test_rsa_verify_should_return_true_if_signature_valid(self):
             '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')

tests/test_api.py

@@ -13,8 +13,9 @@
     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 @@ def default(self, o):
         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()