move user-finding logic out of command

cms/djangoapps/contentstore/management/commands/migrate_to_split.py

@@ -11,6 +11,21 @@ from xmodule.modulestore import InvalidLocationError
 from xmodule.modulestore.django import loc_mapper
 
 
+def user_from_str(s):
+    """
+    Return a user identified by the given string. The string could be an email
+    address, or a stringified integer corresponding to the ID of the user in
+    the database. If no user could be found, a User.DoesNotExist exception
+    will be raised.
+    """
+    try:
+        user_id = int(s)
+    except ValueError:
+        return User.objects.get(email=s)
+    else:
+        return User.objects.get(id=user_id)
+
+
 class Command(BaseCommand):
     help = "Migrate a course from old-Mongo to split-Mongo"
     args = "location email <locator>"
@@ -32,24 +47,10 @@ class Command(BaseCommand):
         except InvalidLocationError:
             raise CommandError("Invalid location string {}".format(args[0]))
 
-        user_id = None
-        email = None
         try:
-            user_id = int(args[1])
-        except ValueError:
-            email = args[1]
-        if user_id:
-            try:
-                user = User.objects.get(pk=user_id)
-            except User.DoesNotExist:
-                raise CommandError("No user exists with ID {}".format(user_id))
-        else:
-            try:
-                user = User.objects.get(email=email)
-            except User.DoesNotExist:
-                raise CommandError("No user exists with email {}".format(email))
-
-        assert user, "User doesn't exist! That shouldn't happen..."
+            user = user_from_str(args[1])
+        except User.DoesNotExist:
+            raise CommandError("No user found identified by {}".format(args[1]))
 
         try:
             locator_string = args[2]

cms/djangoapps/contentstore/management/commands/tests/test_migrate_to_split.py

@@ -29,12 +29,12 @@ class TestArgParsing(unittest.TestCase):
             self.command.handle("foo", "bar")
 
     def test_nonexistant_user_id(self):
-        errstring = "No user exists with ID 99"
+        errstring = "No user found identified by 99"
         with self.assertRaisesRegexp(CommandError, errstring):
             self.command.handle("i4x://org/course/category/name", "99")
 
     def test_nonexistant_user_email(self):
-        errstring = "No user exists with email fake@example.com"
+        errstring = "No user found identified by fake@example.com"
         with self.assertRaisesRegexp(CommandError, errstring):
             self.command.handle("i4x://org/course/category/name", "fake@example.com")
 
@@ -53,11 +53,21 @@ class TestMigrateToSplit(ModuleStoreTestCase):
         self.user = User.objects.create_user(uname, email, password)
         self.course = CourseFactory()
 
-    def test_happy_path(self):
+    def test_happy_path_email(self):
         call_command(
             "migrate_to_split",
             str(self.course.location),
-            self.user.email,
+            str(self.user.email),
+        )
+        locator = loc_mapper().translate_location(self.course.id, self.course.location)
+        course_from_split = modulestore('split').get_course(locator)
+        self.assertIsNotNone(course_from_split)
+
+    def test_happy_path_user_id(self):
+        call_command(
+            "migrate_to_split",
+            str(self.course.location),
+            str(self.user.id),
         )
         locator = loc_mapper().translate_location(self.course.id, self.course.location)
         course_from_split = modulestore('split').get_course(locator)