Add tests for fake followers creation

misago/faker/management/commands/createfakefollowers.py

@@ -14,6 +14,12 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
         total_users = User.objects.count()
+        if total_users < 2:
+            self.stderr.write(
+                "At least two users must exist in the databse in order for "
+                "fake followers creation to be possible.\n"
+            )
+            return
 
         message = "Adding fake followers to %s users...\n"
         self.stdout.write(message % total_users)
@@ -25,20 +31,16 @@ class Command(BaseCommand):
         show_progress(self, processed_count, total_users)
         for user in User.objects.iterator():
             user.followed_by.clear()
-
-            if random.randint(1, 100) > 10:
-                processed_count += 1
-                show_progress(self, processed_count, total_users)
-                continue  # 10% active users
-
-            users_to_add = random.randint(1, total_users / 5)
-            random_queryset = User.objects.exclude(id=user.id).order_by("?")
-            while users_to_add > 0:
-                new_follower = random_queryset[:1][0]
-                if not new_follower.is_following(user):
-                    user.followed_by.add(new_follower)
-                    users_to_add -= 1
-                    total_followers += 1
+            followers_to_create = random.randint(0, total_users - 1)
+            while followers_to_create:
+                # There's 34% chance we'll skip follower creation
+                if random.randint(0, 100) > 34:
+                    new_follower = (
+                        User.objects.exclude(pk=user.pk).order_by("?")[:1].first()
+                    )
+                    if not user.is_following(new_follower):
+                        user.follows.add(new_follower)
+                followers_to_create -= 1
 
             processed_count += 1
             show_progress(self, processed_count, total_users)

misago/faker/tests/test_create_fake_followers_command.py

@@ -0,0 +1,29 @@
+from io import StringIO
+
+from django.core.management import call_command
+
+from ..management.commands import createfakefollowers
+from ..users import get_fake_user
+
+
+def test_management_command_creates_fake_followers_for_two_users(user, other_user):
+    call_command(createfakefollowers.Command(), stdout=StringIO())
+
+
+def test_management_command_creates_fake_followers_for_multiple_users(db, fake):
+    [get_fake_user(fake) for i in range(10)]
+    call_command(createfakefollowers.Command(), stdout=StringIO())
+
+
+def test_management_command_displays_error_if_no_users_exist(db):
+    stderr = StringIO()
+    call_command(createfakefollowers.Command(), stderr=stderr)
+    stderr.seek(0)
+    assert stderr.read()
+
+
+def test_management_command_displays_error_if_only_one_user_exist(user):
+    stderr = StringIO()
+    call_command(createfakefollowers.Command(), stderr=stderr)
+    stderr.seek(0)
+    assert stderr.read()

misago/faker/utils.py

@@ -1,11 +1,12 @@
 from django.db import IntegrityError
+from django.db.transaction import TransactionManagementError
 
 
 def retry_on_db_error(f):
     def wrapper(*args, **kwargs):
         try:
             return f(*args, **kwargs)
-        except IntegrityError:
+        except (IntegrityError, TransactionManagementError):
             return wrapper(*args, **kwargs)
 
     return wrapper