Merge pull request #725 from moggers87/717-get_method_exception

registry.get_method should raise an exception

tests/test_registry.py

@@ -1,6 +1,8 @@
 from django.test import TestCase
 
-from two_factor.plugins.registry import GeneratorMethod, MethodBase, registry
+from two_factor.plugins.registry import (
+    GeneratorMethod, MethodBase, MethodNotFoundError, registry,
+)
 
 
 class FakeMethod(MethodBase):
@@ -43,3 +45,7 @@ def test_unregister_non_registered(self):
 
         registry.unregister('fake-method')
         self.assertEqual(len(registry._methods), expected_length)
+
+    def test_unknown_method(self):
+        with self.assertRaises(MethodNotFoundError):
+            registry.get_method("not-existing-method")

tests/test_views_profile.py

@@ -2,7 +2,7 @@
 from django.test import TestCase, override_settings
 from django.urls import reverse
 
-from two_factor.plugins.registry import registry
+from two_factor.plugins.registry import MethodNotFoundError, registry
 
 from .utils import UserMixin
 
@@ -27,8 +27,10 @@ def test_get_profile_without_phonenumber_plugin_enabled(self):
             app for app in settings.INSTALLED_APPS if app != 'two_factor.plugins.phonenumber']
 
         with override_settings(INSTALLED_APPS=without_phonenumber_plugin):
-            self.assertFalse(registry.get_method('call'))
-            self.assertFalse(registry.get_method('sms'))
+            with self.assertRaises(MethodNotFoundError):
+                registry.get_method('call')
+            with self.assertRaises(MethodNotFoundError):
+                registry.get_method('sms')
 
             response = self.get_profile()
 

two_factor/plugins/phonenumber/utils.py

@@ -2,15 +2,19 @@
 
 import phonenumbers
 
-from two_factor.plugins.registry import registry
+from two_factor.plugins.registry import MethodNotFoundError, registry
 
 phone_mask = re.compile(r'(?<=.{3})[0-9](?=.{2})')
 
 
 def get_available_phone_methods():
     methods = []
     for code in ['sms', 'call']:
-        if method := registry.get_method(code):
+        try:
+            method = registry.get_method(code)
+        except MethodNotFoundError:
+            pass
+        else:
             methods.append(method)
 
     return methods

two_factor/plugins/registry.py

@@ -1,6 +1,15 @@
 from django.utils.translation import gettext_lazy as _
 
 
+class MethodNotFoundError(LookupError):
+    """Registry method was not found
+
+    Check that the appropriate method has been registered
+    """
+    def __init__(self, code, methods):
+        super().__init__(f"{code} not found in {[m.code for m in methods]}")
+
+
 class MethodBase:
     code = None
     verbose_name = None
@@ -90,7 +99,7 @@ def get_method(self, code):
         try:
             return [meth for meth in self._methods if meth.code == code][0]
         except IndexError:
-            return None
+            raise MethodNotFoundError(code, self._methods)
 
     def get_methods(self):
         return self._methods

two_factor/views/core.py

@@ -36,7 +36,7 @@
 
 from two_factor import signals
 from two_factor.plugins.phonenumber.utils import get_available_phone_methods
-from two_factor.plugins.registry import registry
+from two_factor.plugins.registry import MethodNotFoundError, registry
 from two_factor.utils import totp_digits
 from two_factor.views.mixins import OTPRequiredMixin
 
@@ -501,12 +501,13 @@ def get_form_list(self):
             form_list.pop('method', None)
             method_key = available_methods[0].code
             self.storage.validated_step_data['method'] = {'method': method_key}
-        method = self.get_method()
-        if method:
-            form_list.update(method.get_setup_forms(self))
-        else:
+        try:
+            method = self.get_method()
+        except MethodNotFoundError:
             for method in available_methods:
                 form_list.update(method.get_setup_forms(self))
+        else:
+            form_list.update(method.get_setup_forms(self))
         if {'sms', 'call'} & set(form_list.keys()):
             form_list['validation'] = DeviceValidationForm
         return form_list
@@ -549,9 +550,6 @@ def done(self, form_list, **kwargs):
             device = self.get_device()
             device.save()
 
-        else:
-            raise NotImplementedError("Unknown method '%s'" % method.code)
-
         django_otp.login(self.request, device)
         return redirect(self.get_success_url())