refactor phone method availability checks

tests/test_utils.py

@@ -7,10 +7,13 @@
 from phonenumber_field.phonenumber import PhoneNumber
 
 from two_factor.plugins.email.utils import mask_email
+from two_factor.plugins.phonenumber.method import PhoneCallMethod, SMSMethod
 from two_factor.plugins.phonenumber.models import PhoneDevice
 from two_factor.plugins.phonenumber.utils import (
-    backup_phones, format_phone_number, mask_phone_number,
+    backup_phones, format_phone_number, get_available_phone_methods,
+    mask_phone_number,
 )
+from two_factor.plugins.registry import GeneratorMethod, MethodRegistry
 from two_factor.utils import (
     USER_DEFAULT_DEVICE_ATTR_NAME, default_device, get_otpauth_url,
     totp_digits,
@@ -137,11 +140,28 @@ def test_wrong_device_hash(self):
 
 
 class PhoneUtilsTests(UserMixin, TestCase):
+    def test_get_available_phone_methods(self):
+        parameters = [
+            # registered_methods, expected_codes
+            ([GeneratorMethod()], set()),
+            ([GeneratorMethod(), PhoneCallMethod()], {'call'}),
+            ([GeneratorMethod(), PhoneCallMethod(), SMSMethod()], {'call', 'sms'}),
+        ]
+        with mock.patch('two_factor.plugins.phonenumber.utils.registry', new_callable=MethodRegistry) as test_registry:
+            for registered_methods, expected_codes in parameters:
+                with self.subTest(
+                    registered_methods=registered_methods,
+                    expected_codes=expected_codes,
+                ):
+                    test_registry._methods = registered_methods
+                    codes = {method.code for method in get_available_phone_methods()}
+                    self.assertEqual(codes, expected_codes)
+
     def test_backup_phones(self):
         gateway = 'two_factor.gateways.fake.Fake'
         user = self.create_user()
-        user.phonedevice_set.create(name='default', number='+12024561111')
-        backup = user.phonedevice_set.create(name='backup', number='+12024561111')
+        user.phonedevice_set.create(name='default', number='+12024561111', method='call')
+        backup = user.phonedevice_set.create(name='backup', number='+12024561111', method='call')
 
         parameters = [
             # with_gateway, with_user, expected_output

tests/test_views_phone.py

@@ -161,11 +161,13 @@ def setUp(self):
         self.default = self.user.phonedevice_set.create(name='default', method='call', number='+12024561111')
         self.login_user()
 
+    @override_settings(TWO_FACTOR_SMS_GATEWAY='two_factor.gateways.fake.Fake')
     def test_delete(self):
+        self.assertEqual(len(backup_phones(self.user)), 1)
         response = self.client.post(reverse('two_factor:phone_delete',
                                             args=[self.backup.pk]))
         self.assertRedirects(response, resolve_url(settings.LOGIN_REDIRECT_URL))
-        self.assertEqual(list(backup_phones(self.user)), [])
+        self.assertEqual(backup_phones(self.user), [])
 
     def test_cannot_delete_default(self):
         response = self.client.post(reverse('two_factor:phone_delete',

tests/test_views_profile.py

@@ -2,60 +2,42 @@
 from django.test import TestCase, override_settings
 from django.urls import reverse
 
+from two_factor.plugins.registry import registry
+
 from .utils import UserMixin
 
 
+@override_settings(
+    TWO_FACTOR_SMS_GATEWAY='two_factor.gateways.fake.Fake',
+    TWO_FACTOR_CALL_GATEWAY='two_factor.gateways.fake.Fake',
+)
 class ProfileTest(UserMixin, TestCase):
-    PHONENUMBER_PLUGIN_NAME = 'two_factor.plugins.phonenumber'
-    EXPECTED_BASE_CONTEXT_KEYS = {
-        'default_device',
-        'default_device_type',
-        'backup_tokens',
-    }
-    EXPECTED_PHONENUMBER_PLUGIN_ADDITIONAL_KEYS = {
-        'backup_phones',
-        'available_phone_methods',
-    }
-
     def setUp(self):
         super().setUp()
         self.user = self.create_user()
         self.enable_otp()
         self.login_user()
 
-    @classmethod
-    def get_installed_apps_list(cls, with_phone_number_plugin=True):
-        apps = set(settings.INSTALLED_APPS)
-        if with_phone_number_plugin:
-            apps.add(cls.PHONENUMBER_PLUGIN_NAME)
-        else:
-            apps.remove(cls.PHONENUMBER_PLUGIN_NAME)
-        return list(apps)
-
     def get_profile(self):
         url = reverse('two_factor:profile')
         return self.client.get(url)
 
-    def test_get_profile_without_phonenumer_plugin_enabled(self):
-        apps_list = self.get_installed_apps_list(with_phone_number_plugin=False)
-        with override_settings(INSTALLED_APPS=apps_list):
+    def test_get_profile_without_phonenumber_plugin_enabled(self):
+        without_phonenumber_plugin = [
+            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'))
+
             response = self.get_profile()
-            context_keys = set(response.context.keys())
-            self.assertTrue(self.EXPECTED_BASE_CONTEXT_KEYS.issubset(context_keys))
-            # None of the phonenumber related keys are present
-            self.assertTrue(
-                self.EXPECTED_PHONENUMBER_PLUGIN_ADDITIONAL_KEYS.isdisjoint(
-                    context_keys
-                )
-            )
+
+        self.assertTrue(response.context['available_phone_methods'] == [])
 
     def test_get_profile_with_phonenumer_plugin_enabled(self):
-        apps_list = self.get_installed_apps_list(with_phone_number_plugin=True)
-        with override_settings(INSTALLED_APPS=apps_list):
-            response = self.get_profile()
-            context_keys = set(response.context.keys())
-            expected_keys = (
-                self.EXPECTED_BASE_CONTEXT_KEYS
-                | self.EXPECTED_PHONENUMBER_PLUGIN_ADDITIONAL_KEYS
-            )
-            self.assertTrue(expected_keys.issubset(context_keys))
+        self.assertTrue(registry.get_method('call'))
+        self.assertTrue(registry.get_method('sms'))
+
+        response = self.get_profile()
+        available_phone_method_codes = {method.code for method in response.context['available_phone_methods']}
+        self.assertTrue(available_phone_method_codes == {'call', 'sms'})

two_factor/plugins/phonenumber/apps.py

@@ -1,4 +1,4 @@
-from django.apps import AppConfig
+from django.apps import AppConfig, apps
 from django.conf import settings
 from django.test.signals import setting_changed
 
@@ -12,19 +12,21 @@ class TwoFactorPhoneNumberConfig(AppConfig):
     url_prefix = 'phone'
 
     def ready(self):
-        register_methods(self, None, None)
-        setting_changed.connect(register_methods)
+        update_registered_methods(self, None, None)
+        setting_changed.connect(update_registered_methods)
 
 
-def register_methods(sender, setting, value, **kwargs):
+def update_registered_methods(sender, setting, value, **kwargs):
     # This allows for dynamic registration, typically when testing.
     from .method import PhoneCallMethod, SMSMethod
 
-    if getattr(settings, 'TWO_FACTOR_CALL_GATEWAY', None):
+    phone_number_app_installed = apps.is_installed('two_factor.plugins.phonenumber')
+
+    if phone_number_app_installed and getattr(settings, 'TWO_FACTOR_CALL_GATEWAY', None):
         registry.register(PhoneCallMethod())
     else:
         registry.unregister('call')
-    if getattr(settings, 'TWO_FACTOR_SMS_GATEWAY', None):
+    if phone_number_app_installed and getattr(settings, 'TWO_FACTOR_SMS_GATEWAY', None):
         registry.register(SMSMethod())
     else:
         registry.unregister('sms')

two_factor/plugins/phonenumber/forms.py

@@ -15,9 +15,16 @@ class Meta:
         model = PhoneDevice
         fields = ['number', 'method']
 
+    @staticmethod
+    def get_available_choices():
+        choices = []
+        for method in get_available_phone_methods():
+            choices.append((method.code, method.verbose_name))
+        return choices
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.fields['method'].choices = get_available_phone_methods()
+        self.fields['method'].choices = self.get_available_choices()
 
 
 class PhoneNumberForm(forms.ModelForm):

two_factor/plugins/phonenumber/utils.py

@@ -1,33 +1,32 @@
 import re
 
 import phonenumbers
-from django.conf import settings
-from django.utils.translation import gettext_lazy as _
 
-phone_mask = re.compile(r'(?<=.{3})[0-9](?=.{2})')
-
-
-def backup_phones(user):
-    no_gateways = (
-        getattr(settings, 'TWO_FACTOR_CALL_GATEWAY', None) is None
-        and getattr(settings, 'TWO_FACTOR_SMS_GATEWAY', None) is None)
-    no_user = not user or user.is_anonymous
+from two_factor.plugins.registry import registry
 
-    if no_gateways or no_user:
-        from .models import PhoneDevice
-        return PhoneDevice.objects.none()
-    return user.phonedevice_set.filter(name='backup')
+phone_mask = re.compile(r'(?<=.{3})[0-9](?=.{2})')
 
 
 def get_available_phone_methods():
     methods = []
-    if getattr(settings, 'TWO_FACTOR_CALL_GATEWAY', None):
-        methods.append(('call', _('Phone call')))
-    if getattr(settings, 'TWO_FACTOR_SMS_GATEWAY', None):
-        methods.append(('sms', _('Text message')))
+    for code in ['sms', 'call']:
+        if method := registry.get_method(code):
+            methods.append(method)
+
     return methods
 
 
+def backup_phones(user):
+    if not user or user.is_anonymous:
+        return []
+
+    phones = []
+    for method in get_available_phone_methods():
+        phones += list(method.get_devices(user))
+
+    return [phone for phone in phones if phone.name == 'backup']
+
+
 def mask_phone_number(number):
     """
     Masks a phone number, only first 3 and last 2 digits visible.

two_factor/views/profile.py

@@ -1,4 +1,3 @@
-from django.apps.registry import apps
 from django.conf import settings
 from django.contrib.auth.decorators import login_required
 from django.shortcuts import redirect, resolve_url
@@ -30,24 +29,22 @@ class ProfileView(TemplateView):
     template_name = 'two_factor/profile/profile.html'
 
     def get_context_data(self, **kwargs):
+        user = self.request.user
+
         try:
-            backup_tokens = self.request.user.staticdevice_set.all()[0].token_set.count()
+            backup_tokens = user.staticdevice_set.all()[0].token_set.count()
 
         except Exception:
             backup_tokens = 0
 
         context = {
-            'default_device': default_device(self.request.user),
-            'default_device_type': default_device(self.request.user).__class__.__name__,
+            'default_device': default_device(user),
+            'default_device_type': default_device(user).__class__.__name__,
             'backup_tokens': backup_tokens,
+            'backup_phones': backup_phones(user),
+            'available_phone_methods': get_available_phone_methods(),
         }
 
-        if (apps.is_installed('two_factor.plugins.phonenumber')):
-            context.update({
-                'backup_phones': backup_phones(self.request.user),
-                'available_phone_methods': get_available_phone_methods(),
-            })
-
         return context