Add a util script to create large amounts of random data

bakerydemo/base/management/commands/create_random_data.py

@@ -0,0 +1,157 @@
+from django.core.management.base import BaseCommand
+
+from bakerydemo.breads.models import BreadPage, BreadsIndexPage, BreadType, Country, BreadIngredient
+from django.conf import settings
+from django.utils import lorem_ipsum, timezone
+from django.utils.text import slugify
+from wagtail.core.rich_text import RichText
+from bakerydemo.locations.models import LocationPage, LocationsIndexPage
+from bakerydemo.blog.models import BlogPage, BlogIndexPage
+import random
+from wagtail.images.models import Image
+from pathlib import Path
+from willow.image import Image as WillowImage
+from django.conf import settings
+from bakerydemo.base.models import People, FooterText, StandardPage, HomePage
+
+FIXTURE_MEDIA_DIR = Path(settings.PROJECT_DIR) / "base/fixtures/media/original_images"
+
+class Command(BaseCommand):
+    help = "Creates random data. Useful for performance or load testing."
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            'page_count', type=int,
+            help='How many pages of each type to create',
+        )
+        parser.add_argument(
+            'snippet_count', type=int,
+            help='How many snippets of each type to create',
+        )
+        parser.add_argument(
+            'image_count', type=int,
+            help='How many images to create',
+        )
+
+    def fake_stream_field(self):
+        return [
+            ('paragraph_block', RichText("\n".join(lorem_ipsum.paragraphs(5))))
+        ]
+
+    def get_random_model(self, model):
+        return model.objects.order_by('?').first()
+
+    def make_title(self):
+        return lorem_ipsum.words(4, common=False)
+
+    def create_pages(self, page_count):
+        self.stdout.write("Creating bread pages...")
+        breads_index = BreadsIndexPage.objects.live().first()
+        for _ in range(page_count):
+            title = self.make_title()
+            breads_index.add_child(instance=BreadPage(
+                title=title,
+                slug=slugify(title),
+                introduction=lorem_ipsum.paragraph(),
+                bread_type=self.get_random_model(BreadType),
+                body=self.fake_stream_field(),
+                origin=self.get_random_model(Country),
+                image=self.get_random_model(Image)
+            ))
+
+        self.stdout.write("Creating location pages...")
+        locations_index = LocationsIndexPage.objects.live().first()
+        for _ in range(page_count):
+            title = self.make_title()
+            locations_index.add_child(instance=LocationPage(
+                title=title,
+                slug=slugify(title),
+                introduction=lorem_ipsum.paragraph(),
+                image=self.get_random_model(Image),
+                address=lorem_ipsum.paragraph(),
+                body=self.fake_stream_field(),
+                lat_long='64.144367, -21.939182'
+            ))
+
+        self.stdout.write("Creating blog pages...")
+        blog_index = BlogIndexPage.objects.live().first()
+        for _ in range(page_count):
+            title = self.make_title()
+            blog_index.add_child(instance=BlogPage(
+                title=title,
+                slug=slugify(title),
+                introduction=lorem_ipsum.paragraph(),
+                body=self.fake_stream_field(),
+                subtitle=lorem_ipsum.words(10, common=False),
+                date_published=timezone.now()
+            ))
+
+        self.stdout.write("Creating standard pages...")
+        homepage = HomePage.objects.live().first()
+        title = self.make_title()
+        # Nest the standard pages under a top level one
+        top_level_page = homepage.add_child(instance=StandardPage(
+            title=title,
+            slug=slugify(title),
+            introduction=lorem_ipsum.paragraph(),
+            image=self.get_random_model(Image),
+            body=self.fake_stream_field(),
+        ))
+        for _ in range(page_count):
+            title = self.make_title()
+            top_level_page.add_child(instance=StandardPage(
+                title=title,
+                slug=slugify(title),
+                introduction=lorem_ipsum.paragraph(),
+                image=self.get_random_model(Image),
+                body=self.fake_stream_field(),
+            ))
+
+    def create_snippets(self, snippet_count):
+        self.stdout.write("Creating countries...")
+        for _ in range(snippet_count):
+            Country.objects.create(title=self.make_title())
+
+        self.stdout.write("Creating bread ingredients...")
+        for _ in range(snippet_count):
+            BreadIngredient.objects.create(name=self.make_title())
+
+        self.stdout.write("Creating bread types...")
+        for _ in range(snippet_count):
+            BreadType.objects.create(title=self.make_title())
+
+        self.stdout.write("Creating people...")
+        for _ in range(snippet_count):
+            People.objects.create(
+                first_name=lorem_ipsum.words(1, common=False),
+                last_name=lorem_ipsum.words(1, common=False),
+                job_title=lorem_ipsum.words(1, common=False),
+                image=self.get_random_model(Image)
+            )
+
+        self.stdout.write("Creating footer text...")
+        for _ in range(snippet_count):
+            FooterText.objects.create(body=self.fake_stream_field())
+
+    def create_images(self, image_count):
+        image_files = list(FIXTURE_MEDIA_DIR.iterdir())
+
+        self.stdout.write("Creating images...")
+        for _ in range(image_count):
+            random_image = random.choice(image_files)
+            with random_image.open(mode="rb") as image_file:
+                willow_image = WillowImage.open(image_file)
+                width, height = willow_image.get_size()
+                image = Image.objects.create(
+                    title=self.make_title(),
+                    width=width,
+                    height=height,
+                    file_size=random_image.stat().st_size,
+                )
+                image_file.seek(0)
+                image.file.save(random_image.name, image_file)
+
+    def handle(self, **options):
+        self.create_images(options['image_count'])
+        self.create_snippets(options['snippet_count'])
+        self.create_pages(options['page_count'])