Skip to content

Commit

Permalink
fix: issue with PhoneDevice being imported when not enabled
Browse files Browse the repository at this point in the history
  • Loading branch information
dhruuuuuv authored Aug 28, 2023
1 parent 08ef1b8 commit 6e7eca4
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 3 deletions.
61 changes: 61 additions & 0 deletions tests/test_views_profile.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
from django.conf import settings
from django.test import TestCase, override_settings
from django.urls import reverse

from .utils import UserMixin


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):
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
)
)

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))
14 changes: 11 additions & 3 deletions two_factor/views/profile.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
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
Expand Down Expand Up @@ -31,17 +32,24 @@ class ProfileView(TemplateView):
def get_context_data(self, **kwargs):
try:
backup_tokens = self.request.user.staticdevice_set.all()[0].token_set.count()

except Exception:
backup_tokens = 0

return {
context = {
'default_device': default_device(self.request.user),
'default_device_type': default_device(self.request.user).__class__.__name__,
'backup_phones': backup_phones(self.request.user),
'backup_tokens': backup_tokens,
'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


@class_view_decorator(never_cache)
class DisableView(FormView):
Expand Down

0 comments on commit 6e7eca4

Please sign in to comment.