From ce0fc3804d146136b6e1c63dc1342b566f8b0b2d Mon Sep 17 00:00:00 2001 From: Awais Jibran Date: Thu, 4 Mar 2021 15:16:54 +0500 Subject: [PATCH] refactor: use q objects when fetching user using an identifier #26683 (#26823) --- .../management/commands/cert_whitelist.py | 12 ++++----- .../commands/tests/test_cert_whitelist.py | 25 ++++++++++++++++--- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/lms/djangoapps/certificates/management/commands/cert_whitelist.py b/lms/djangoapps/certificates/management/commands/cert_whitelist.py index 410b11fc51..c019048c9f 100644 --- a/lms/djangoapps/certificates/management/commands/cert_whitelist.py +++ b/lms/djangoapps/certificates/management/commands/cert_whitelist.py @@ -6,6 +6,7 @@ user/course from django.contrib.auth import get_user_model from django.core.management.base import BaseCommand, CommandError +from django.db.models import Q from opaque_keys.edx.keys import CourseKey from lms.djangoapps.certificates.models import CertificateWhitelist @@ -17,11 +18,9 @@ def get_user_from_identifier(identifier): """ This function takes the string identifier and fetch relevant user object from database """ - identifier = identifier.strip() - if '@' in identifier: - user = User.objects.get(email=identifier) - else: - user = User.objects.get(username=identifier) + user = User.objects.filter(Q(username=identifier) | Q(email=identifier)).first() + if not user: + raise CommandError("User {} does not exist.".format(identifier)) return user @@ -103,7 +102,8 @@ class Command(BaseCommand): add_to_whitelist = True if options['add'] else False # pylint: disable=simplifiable-if-expression users_list = user_str.split(",") for username in users_list: - if username.strip(): + username = username.strip() + if username: update_user_whitelist(username, add=add_to_whitelist) whitelist = CertificateWhitelist.objects.filter(course_id=course) diff --git a/lms/djangoapps/certificates/management/commands/tests/test_cert_whitelist.py b/lms/djangoapps/certificates/management/commands/tests/test_cert_whitelist.py index db5f80690e..bc76c45d38 100644 --- a/lms/djangoapps/certificates/management/commands/tests/test_cert_whitelist.py +++ b/lms/djangoapps/certificates/management/commands/tests/test_cert_whitelist.py @@ -1,10 +1,10 @@ """ Extremely basic tests for the cert_whitelist command """ - - +import ddt import pytest -from django.core.management import call_command +from django.core.management import call_command, CommandError +from xmodule.modulestore.tests.django_utils import ModuleStoreTestCase def test_cert_whitelist_help(capsys): @@ -16,3 +16,22 @@ def test_cert_whitelist_help(capsys): out, err = capsys.readouterr() # pylint: disable=unused-variable assert "COURSE_ID" in out + + +@ddt.ddt +class CertWhitelistGenerationTests(ModuleStoreTestCase): + """ + Tests for the cert_whitelist management command. + """ + @ddt.data( + 'jeo', + 'jeo@edx.org', + 'robo_user, jenny, tom, jerry' + ) + def test_user_not_found(self, user_identifier): + """ + Basic test to see if the command will raise user not found error. + """ + error_identifier = user_identifier.split(',')[0].strip() + with pytest.raises(CommandError, match=f"User {error_identifier} does not exist"): + call_command('cert_whitelist', '--add', user_identifier, '-c', 'MITx/6.002x/2012_Fall')