GlitchTip/issues/management/commands/make_sample_issues.py

import random

from django.core.management.base import BaseCommand
from model_bakery import baker
from model_bakery.random_gen import gen_json, gen_slug

from events.test_data import event_generator
from events.views import EventStoreAPIView
from projects.models import Project

baker.generators.add("organizations.fields.SlugField", gen_slug)
baker.generators.add("django.db.models.JSONField", gen_json)


class Command(BaseCommand):
    help = "Create sample issues and events for dev and demonstration purposes"

    def add_arguments(self, parser):
        parser.add_argument("quantity", nargs="?", type=int)
        parser.add_argument(
            "--only-real", action="store_true", help="Only include real sample events",
        )
        parser.add_argument(
            "--only-fake",
            action="store_true",
            help="Only include faked generated events",
        )

    def generate_real_event(self, project, unique_issue=False):
        """ Generate an event based on real sample data """
        data = event_generator.generate_random_event(unique_issue)
        serializer = EventStoreAPIView().get_event_serializer_class(data)(
            data=data, context={"project": project}
        )
        serializer.is_valid()
        serializer.save()

    def handle(self, *args, **options):
        project = Project.objects.first()
        if not project:
            project = baker.make("projects.Project")
        if options["quantity"] is None:
            options["quantity"] = 1
        quantity = options["quantity"]

        only_real = options["only_real"]
        only_fake = options["only_fake"]

        if only_real:
            for _ in range(quantity):
                self.generate_real_event(project)
        elif only_fake:
            baker.make("events.Event", issue__project=project, _quantity=quantity)
        else:
            for _ in range(quantity):
                if random.choice([0, 1]):  # nosec
                    baker.make("events.Event", issue__project=project)
                else:
                    self.generate_real_event(project)

        self.stdout.write(
            self.style.SUCCESS('Successfully created "%s" events' % quantity)
        )