feat: Add command to load test tracking events

platform_plugin_aspects/management/commands/load_test_tracking_events.py

@@ -0,0 +1,197 @@
+"""
+Generates tracking events by creating test users and fake activity.
+
+This should never be run on a production server as it will generate a lot of
+bad data. It is entirely for benchmarking purposes in load test environments.
+It is also fragile due to reaching into the edx-platform testing internals.
+"""
+
+import logging
+import uuid
+from datetime import datetime, timedelta
+from random import choice
+from textwrap import dedent
+from time import sleep
+from typing import Any
+
+from django.contrib.auth.models import User
+from django.core.management.base import BaseCommand, CommandError
+
+try:
+    from cms.djangoapps.contentstore.views.course import create_new_course_in_store
+    from common.djangoapps.student.helpers import do_create_account
+    from common.djangoapps.student.models.course_enrollment import CourseEnrollment
+    from openedx.core.djangoapps.user_authn.views.registration_form import (
+        AccountCreationForm,
+    )
+    from xmodule.modulestore import ModuleStoreEnum
+
+    RUNNING_IN_PLATFORM = True
+except ImportError:
+    RUNNING_IN_PLATFORM = False
+
+log = logging.getLogger(__name__)
+
+
+class LoadTest:
+    """
+    Base class for setting up and sending events.
+    """
+
+    course = None
+    instructor = None
+    users = []
+    sent_event_count = 0
+
+    def __init__(self, num_users: int, username_prefix: str):
+        course_shortname = str(uuid.uuid4())[:6]
+        self.instructor = self.create_user(
+            username=f"instructor_{course_shortname}",
+            name="Instructor",
+            password="aspects",
+            email=f"instructor_{course_shortname}@openedx.invalid",
+        )
+
+        start_date = datetime.now() - timedelta(days=7)
+
+        fields = {"start": start_date, "display_name": f"Course {course_shortname}"}
+
+        log.info(
+            f"""Creating course:
+        Instructor: {self.instructor.id}
+        Org: "OEX"
+        Number: "{course_shortname}"
+        Run: "2024-1"
+        Fields: {fields}
+        """
+        )
+
+        self.course = create_new_course_in_store(
+            ModuleStoreEnum.Type.split,
+            self.instructor,
+            "OEX",
+            course_shortname,
+            "2024-1",
+            fields,
+        )
+
+        log.info(f"Created course {self.course.id}")
+        self.create_and_enroll_learners(num_users, username_prefix)
+
+    def create_and_enroll_learners(self, num_users, username_prefix):
+        log.info(f"Creating {num_users} users prefixed with {username_prefix}.")
+
+        for _ in range(num_users):
+            user_short_name = str(uuid.uuid4())[:6]
+            u = self.create_user(
+                username=f"{username_prefix}_{user_short_name}",
+                name=f"Learner {user_short_name}",
+                password="aspects",
+                email=f"{user_short_name}@openedx.invalid",
+            )
+            self.users.append(u)
+            e = CourseEnrollment.get_or_create_enrollment(
+                user=u, course_key=self.course.id
+            )
+            e.is_active = True
+            e.save()
+
+    def create_user(self, **user_data):
+        account_creation_form = AccountCreationForm(data=user_data, tos_required=False)
+
+        user, _, _ = do_create_account(account_creation_form)
+        user.is_active = True
+        user.save()
+        return user
+
+    def trigger_events(
+        self, num_events: int, sleep_time: float, run_until_killed: bool
+    ) -> None:
+        if run_until_killed:
+            log.info(f"Creating events until killed with {sleep_time} between!")
+            while True:
+                self.trigger_event_and_sleep(sleep_time)
+        else:
+            log.info(f"Creating events {num_events} with {sleep_time} between!")
+            for _ in range(num_events):
+                self.trigger_event_and_sleep(sleep_time)
+
+    def trigger_event_and_sleep(self, sleep_time: float) -> None:
+        user = choice(self.users)
+        log.info(f"Triggering event for user {user.username}.")
+        e = CourseEnrollment.get_or_create_enrollment(
+            user=user, course_key=self.course.id
+        )
+
+        if e.is_active:
+            e.unenroll(user, self.course.id)
+        else:
+            e.enroll(user, self.course.id)
+
+        self.sent_event_count += 1
+        sleep(sleep_time)
+
+
+class Command(BaseCommand):
+    """
+    Create tracking log events for load testing purposes.
+
+    Example:
+    tutor local run lms ./manage.py lms load_test_tracking_events --sleep_time 0
+    """
+
+    help = dedent(__doc__).strip()
+
+    def add_arguments(self, parser: Any) -> None:
+        parser.add_argument(
+            "--num_users",
+            type=int,
+            default=10,
+            help="The number of users to create. All events will be generated for these learners.",
+        )
+        parser.add_argument(
+            "--username_prefix",
+            type=str,
+            default="lt_",
+            help="Prefix for the generated user names.",
+        )
+        parser.add_argument(
+            "--num_events",
+            type=int,
+            default=10,
+            help="The number of events to generate. This is ignored if --run_until_killed is set.",
+        )
+        parser.add_argument(
+            "--run_until_killed",
+            action="store_true",
+            default=False,
+            help="If this is set, the process will run endlessly until killed.",
+        )
+        parser.add_argument(
+            "--sleep_time",
+            type=float,
+            default=0.75,
+            help="Fractional number of seconds to sleep between sending events.",
+        )
+
+    def handle(self, *args, **options):
+        """
+        Creates users and triggers events for them as configured above.
+        """
+        if not RUNNING_IN_PLATFORM:
+            raise CommandError("This command must be run in the Open edX LMS or CMS.")
+
+        start = datetime.now()
+        lt = LoadTest(options["num_users"], options["username_prefix"])
+
+        try:
+            lt.trigger_events(
+                options["num_events"],
+                options["sleep_time"],
+                options["run_until_killed"],
+            )
+        except KeyboardInterrupt:
+            log.warning("Killed by keyboard, finishing.")
+
+        end = datetime.now()
+        log.info(f"Sent {lt.sent_event_count} events in {end - start}.")