make_sample_issues.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. import random
  2. from django.core.management.base import BaseCommand
  3. from model_bakery import baker
  4. from model_bakery.random_gen import gen_json, gen_slug
  5. from events.test_data import event_generator
  6. from events.views import EventStoreAPIView
  7. from projects.models import Project
  8. baker.generators.add("organizations.fields.SlugField", gen_slug)
  9. baker.generators.add("django.db.models.JSONField", gen_json)
  10. class Command(BaseCommand):
  11. help = "Create sample issues and events for dev and demonstration purposes"
  12. def add_arguments(self, parser):
  13. parser.add_argument("quantity", nargs="?", type=int)
  14. parser.add_argument(
  15. "--only-real",
  16. action="store_true",
  17. help="Only include real sample events",
  18. )
  19. parser.add_argument(
  20. "--only-fake",
  21. action="store_true",
  22. help="Only include faked generated events",
  23. )
  24. def generate_real_event(self, project, unique_issue=False):
  25. """Generate an event based on real sample data"""
  26. data = event_generator.generate_random_event(unique_issue)
  27. project.release_id = None
  28. project.environment_id = None
  29. serializer = EventStoreAPIView().get_event_serializer_class(data)(
  30. data=data, context={"project": project}
  31. )
  32. serializer.is_valid()
  33. serializer.save()
  34. def handle(self, *args, **options):
  35. project = Project.objects.first()
  36. if not project:
  37. project = baker.make("projects.Project")
  38. if options["quantity"] is None:
  39. options["quantity"] = 1
  40. quantity = options["quantity"]
  41. only_real = options["only_real"]
  42. only_fake = options["only_fake"]
  43. if only_real:
  44. for _ in range(quantity):
  45. self.generate_real_event(project)
  46. elif only_fake:
  47. baker.make("events.Event", issue__project=project, _quantity=quantity)
  48. else:
  49. for _ in range(quantity):
  50. if random.choice([0, 1]): # nosec
  51. baker.make("events.Event", issue__project=project)
  52. else:
  53. self.generate_real_event(project)
  54. self.stdout.write(
  55. self.style.SUCCESS('Successfully created "%s" events' % quantity)
  56. )