move get_key to jwk, return typeerror in register_key; updated tests

Author: Marko Mrdjenovic

jose/constants.py

@@ -35,25 +35,13 @@ class Algorithms(object):
 
     KEYS = {}
 
-    def get_key(self, algorithm):
-        from jose.jwk import HMACKey, RSAKey, ECKey
-        if algorithm in self.KEYS:
-            return self.KEYS[algorithm]
-        elif algorithm in self.HMAC:
-            return HMACKey
-        elif algorithm in self.RSA:
-            return RSAKey
-        elif algorithm in self.EC:
-            return ECKey
-        return None
-
     def register_key(self, algorithm, key_class):
         from jose.jwk import Key
-        if issubclass(key_class, Key):
-            self.KEYS[algorithm] = key_class
-            self.SUPPORTED.add(algorithm)
-            return True
-        else:
-            return False
+        if not issubclass(key_class, Key):
+            raise TypeError("Key class not a subclass of jwk.Key")
+        self.KEYS[algorithm] = key_class
+        self.SUPPORTED.add(algorithm)
+        return True
+
 
 ALGORITHMS = Algorithms()

jose/jwk.py

@@ -47,6 +47,18 @@ def base64_to_long(data):
     return int_arr_to_long(struct.unpack('%sB' % len(_d), _d))
 
 
+def get_key(algorithm):
+    if algorithm in ALGORITHMS.KEYS:
+        return ALGORITHMS.KEYS[algorithm]
+    elif algorithm in ALGORITHMS.HMAC:
+        return HMACKey
+    elif algorithm in ALGORITHMS.RSA:
+        return RSAKey
+    elif algorithm in ALGORITHMS.EC:
+        return ECKey
+    return None
+
+
 def construct(key_data, algorithm=None):
     """
     Construct a Key object for the given algorithm with the given
@@ -60,7 +72,7 @@ def construct(key_data, algorithm=None):
     if not algorithm:
         raise JWKError('Unable to find a algorithm for key: %s' % key_data)
 
-    key_class = ALGORITHMS.get_key(algorithm)
+    key_class = get_key(algorithm)
     if not key_class:
         raise JWKError('Unable to find a algorithm for key: %s' % key_data)
     return key_class(key_data, algorithm)

tests/algorithms/test_base.py

@@ -23,15 +23,10 @@ def test_verify_is_interface(self, alg):
 
 class TestAlgorithms:
 
-    def test_get_key(self):
-        assert ALGORITHMS.get_key("HS256") == HMACKey
-        assert ALGORITHMS.get_key("RS256") == RSAKey
-        assert ALGORITHMS.get_key("ES256") == ECKey
-
-        assert ALGORITHMS.get_key("NONEXISTENT") == None
-    
     def test_register_key(self):
-        assert ALGORITHMS.register_key("ALG", Key) == True
-        assert ALGORITHMS.get_key("ALG") == Key
-
-        assert ALGORITHMS.register_key("ALG", object) == False
+        assert ALGORITHMS.register_key("ALG", Key)
+        from jose.jwk import get_key
+        assert get_key("ALG") == Key
+    
+        with pytest.raises(TypeError):
+            assert ALGORITHMS.register_key("ALG", object)

tests/test_jwk.py

@@ -115,3 +115,10 @@ def test_construct_from_jwk_missing_alg(self):
 
         with pytest.raises(JWKError):
             key = jwk.construct("key", algorithm="NONEXISTENT")
+
+    def test_get_key(self):
+        assert jwk.get_key("HS256") == jwk.HMACKey
+        assert jwk.get_key("RS256") == jwk.RSAKey
+        assert jwk.get_key("ES256") == jwk.ECKey
+
+        assert jwk.get_key("NONEXISTENT") == None