Management command for back-populating the 'created_on_site' in UserAttribute model.

common/djangoapps/student/management/commands/populate_created_on_site_user_attribute.py

+"""
+Command to back-populate domain of the site the user account was created on.
+"""
+from django.contrib.auth.models import User
+from django.contrib.sites.models import Site
+from django.core.management.base import BaseCommand, CommandError
+
+from student.models import UserAttribute, Registration
+CREATED_ON_SITE = 'created_on_site'
+
+
+class Command(BaseCommand):
+    """
+    This command back-populates domain of the site the user account was created on.
+    """
+    help = """./manage.py lms populate_created_on_site_user_attribute --users <user_id1>,<user_id2>...
+           '--activation-keys <key1>,<key2>... --site-domain <site_domain> --settings=devstack"""
+
+    def add_arguments(self, parser):
+        """
+        Add arguments to the command parser.
+        """
+        parser.add_argument(
+            '--users',
+            help='Enter comma-separated user ids.',
+            default='',
+            type=str
+        )
+        parser.add_argument(
+            '--activation-keys',
+            help='Enter comma-separated activation keys.',
+            default='',
+            type=str
+        )
+        parser.add_argument(
+            '--site-domain',
+            help='Enter an existing site domain.',
+        )
+
+    def handle(self, *args, **options):
+        site_domain = options['site_domain']
+        user_ids = options['users'].split(',') if options['users'] else []
+        activation_keys = options['activation_keys'].split(',') if options['activation_keys'] else []
+
+        if not site_domain:
+            raise CommandError('You must provide site-domain argument.')
+        if not user_ids and not activation_keys:
+            raise CommandError('You must provide user ids or activation keys.')
+
+        try:
+            Site.objects.get(domain__exact=site_domain)
+        except Site.DoesNotExist:
+            question = "The site you specified is not configured as a Site in the system. " \
+                       "Are you sure you want to continue? (y/n):"
+            if str(raw_input(question)).lower().strip()[0] != 'y':
+                return
+
+        for user_id in user_ids:
+            try:
+                user = User.objects.get(id=user_id)
+                if UserAttribute.get_user_attribute(user, CREATED_ON_SITE):
+                    self.stdout.write("created_on_site attribute already exists for user id: {id}".format(id=user_id))
+                else:
+                    UserAttribute.set_user_attribute(user, CREATED_ON_SITE, site_domain)
+            except User.DoesNotExist:
+                self.stdout.write("This user id [{id}] does not exist in the system.".format(id=user_id))
+
+        for key in activation_keys:
+            try:
+                user = Registration.objects.get(activation_key=key).user
+                if UserAttribute.get_user_attribute(user, CREATED_ON_SITE):
+                    self.stdout.write("created_on_site attribute already exists for user id: {id}".format(id=user.id))
+                else:
+                    UserAttribute.set_user_attribute(user, CREATED_ON_SITE, site_domain)
+            except Registration.DoesNotExist:
+                self.stdout.write("This activation key [{key}] does not exist in the system.".format(key=key))

common/djangoapps/student/management/tests/test_populate_created_on_site_user_attribute.py

+"""
+Unittests for populate_created_on_site_user_attribute management command.
+"""
+import ddt
+import mock
+from django.test import TestCase
+from django.contrib.auth.models import User
+from django.core.management import call_command, CommandError
+
+from student.models import Registration, UserAttribute
+from student.tests.factories import UserFactory
+from openedx.core.djangoapps.site_configuration.tests.mixins import SiteMixin
+CREATED_ON_SITE = 'created_on_site'
+
+
+@ddt.ddt
+class TestPopulateUserAttribute(SiteMixin, TestCase):
+    """
+    Test populate_created_on_site_user_attribute management command.
+    """
+
+    def setUp(self):
+        super(TestPopulateUserAttribute, self).setUp()
+
+        self._create_sample_data()
+        self.users = User.objects.all()
+        self.registered_users = Registration.objects.all()
+        self.user_ids = ','.join([str(user.id) for user in self.users])
+        self.activation_keys = ','.join([registered_user.activation_key for registered_user in self.registered_users])
+
+    def _create_sample_data(self):
+        """
+        Creates the users and register them.
+        """
+        for __ in range(3):
+            Registration().register(UserFactory.create())
+
+    def test_command_by_user_ids(self):
+        """
+        Test population of created_on_site attribute by user ids.
+        """
+        call_command(
+            "populate_created_on_site_user_attribute",
+            "--users", self.user_ids,
+            "--site-domain", self.site.domain
+        )
+
+        for user in self.users:
+            self.assertEqual(UserAttribute.get_user_attribute(user, CREATED_ON_SITE), self.site.domain)
+
+        # Populate 'created_on_site' attribute with different site domain
+        call_command(
+            "populate_created_on_site_user_attribute",
+            "--users", self.user_ids,
+            "--site-domain", self.site_other.domain
+        )
+
+        for user in self.users:
+            # 'created_on_site' attribute already exists. Attribute's value will not change
+            self.assertNotEqual(UserAttribute.get_user_attribute(user, CREATED_ON_SITE), self.site_other.domain)
+
+    def test_command_by_activation_keys(self):
+        """
+        Test population of created_on_site attribute by activation keys.
+        """
+        call_command(
+            "populate_created_on_site_user_attribute",
+            "--activation-keys", self.activation_keys,
+            "--site-domain", self.site.domain
+        )
+
+        for register_user in self.registered_users:
+            self.assertEqual(UserAttribute.get_user_attribute(register_user.user, CREATED_ON_SITE), self.site.domain)
+
+        # Populate 'created_on_site' attribute with different site domain
+        call_command(
+            "populate_created_on_site_user_attribute",
+            "--activation-keys", self.activation_keys,
+            "--site-domain", self.site_other.domain
+        )
+
+        for register_user in self.registered_users:
+            # 'created_on_site' attribute already exists. Attribute's value will not change
+            self.assertNotEqual(
+                UserAttribute.get_user_attribute(register_user.user, CREATED_ON_SITE),
+                self.site_other.domain
+            )
+
+    def test_command_with_incomplete_argument(self):
+        """
+        Test management command raises CommandError without '--users' and '--activation_keys' arguments.
+        """
+        with self.assertRaises(CommandError):
+            call_command(
+                "populate_created_on_site_user_attribute",
+                "--site-domain", self.site.domain
+            )
+
+    def test_command_with_invalid_arguments(self):
+        """
+        Test management command with invalid user ids and activation keys.
+        """
+        user = self.users[0]
+        call_command(
+            "populate_created_on_site_user_attribute",
+            "--users", '9{id}'.format(id=user.id),  # invalid id
+            "--site-domain", self.site.domain
+        )
+        self.assertIsNone(UserAttribute.get_user_attribute(user, CREATED_ON_SITE))
+
+        register_user = self.registered_users[0]
+        call_command(
+            "populate_created_on_site_user_attribute",
+            "--activation-keys", "invalid-{key}".format(key=register_user.activation_key),  # invalid key
+            "--site-domain", self.site.domain
+        )
+        self.assertIsNone(UserAttribute.get_user_attribute(register_user.user, CREATED_ON_SITE))
+
+    def test_command_without_site_domain(self):
+        """
+        Test management command raises CommandError without '--site-domain' argument.
+        """
+        with self.assertRaises(CommandError):
+            call_command(
+                "populate_created_on_site_user_attribute",
+                "--user", self.user_ids,
+                "--activation-keys", self.activation_keys
+            )
+
+    @ddt.data('y', 'n')
+    def test_with_invalid_site_domain(self, populate):
+        """
+        Test management command with invalid site domain.
+        """
+        fake_site_domain = 'fake-site-domain'
+        with mock.patch('__builtin__.raw_input', return_value=populate):
+            call_command(
+                "populate_created_on_site_user_attribute",
+                "--users", self.user_ids,
+                "--site-domain", fake_site_domain
+            )
+
+        for user in self.users:
+            if populate == 'y':
+                self.assertEqual(UserAttribute.get_user_attribute(user, CREATED_ON_SITE), fake_site_domain)
+            else:
+                self.assertIsNone(UserAttribute.get_user_attribute(user, CREATED_ON_SITE))