Custom Django command to seed database

You can execute it with the following command: python manage.py seed_db

django_graphql_playground/apps/core/management/commands/seed_db.py

@@ -0,0 +1,50 @@
+from django.contrib.auth import get_user_model
+from django.core.management import BaseCommand
+from django.core.management import CommandError
+
+from django_graphql_playground.apps.core.models import Category
+from django_graphql_playground.apps.core.models import Ingredient
+
+
+class Command(BaseCommand):
+    help = "Seed database with sample data"
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--create-super-user", action="store_true", dest="create_super_user", help="Create default super user"
+        )
+        parser.add_argument(
+            "-u", dest="admin_username", type=str, default="admin", help="Super user username. Defaults to: admin "
+        )
+        parser.add_argument(
+            "-p",
+            dest="admin_password",
+            type=str,
+            default="Asd123!.",
+            help="Super user password. Defaults to: Asd123!. ",
+        )
+
+    def handle(self, *args, **options):
+        self.create_super_user = options["create_super_user"]
+        self.admin_username = options["admin_username"].strip()
+        self.admin_password = options["admin_password"].strip()
+
+        if self.create_super_user == True:
+            user_model = get_user_model()
+            if user_model.objects.filter(username=self.admin_username).count() == 0:
+                self.stdout.write(f"Creating with username {self.admin_username} and password {self.admin_password}")
+                get_user_model().objects.create_superuser(self.admin_username, None, self.admin_password)
+            else:
+                raise CommandError(f"Super user {self.admin_username} already exists")
+
+        if Category.objects.all().count() == 0 and Ingredient.objects.all().count() == 0:
+            self.stdout.write(f"Creating categories")
+            category_one = Category.objects.create(name="Category XPTO")
+            category_two = Category.objects.create(name="Category QWERTY")
+            self.stdout.write(f"Creating ingredients")
+            Ingredient.objects.create(name="Salt", notes=None, category=category_one)
+            Ingredient.objects.create(name="Oregano", notes=None, category=category_one)
+            Ingredient.objects.create(name="Cumin", notes=None, category=category_one)
+            Ingredient.objects.create(name="Chive", notes=None, category=category_two)
+        else:
+            self.stdout.write(f"There are categories and ingredients already")

tests/integration/apps/core/management/commands/test_seed_db.py

@@ -0,0 +1,64 @@
+from io import StringIO
+
+import pytest
+from django.contrib.auth import get_user_model
+from django.contrib.auth.models import User
+from django.core.management import CommandError
+from django.core.management import call_command
+
+from django_graphql_playground.apps.core.models import Category
+from django_graphql_playground.apps.core.models import Ingredient
+
+
+@pytest.mark.django_db
+def test_should_seed_db_without_creating_super_user():
+    out = StringIO()
+
+    call_command("seed_db", stdout=out)
+
+    assert out.getvalue() == "Creating categories\nCreating ingredients\n"
+    assert Category.objects.all().count() == 2
+    assert Ingredient.objects.all().count() == 4
+
+
+@pytest.mark.django_db
+def test_should_seed_db_with_creating_super_user_given_extra_option():
+    out = StringIO()
+
+    call_command("seed_db", "--create-super-user", stdout=out)
+
+    assert (
+        out.getvalue()
+        == "Creating with username admin and password Asd123!.\nCreating categories\nCreating ingredients\n"
+    )
+    assert User.objects.filter(username="admin").count() == 1
+    assert Category.objects.all().count() == 2
+    assert Ingredient.objects.all().count() == 4
+
+
+@pytest.mark.django_db
+def test_should_seed_db_with_custom_super_user_given_extra_options():
+    out = StringIO()
+    custom_username = "xpto"
+
+    call_command("seed_db", "--create-super-user", f"-u {custom_username}", stdout=out)
+
+    assert (
+        out.getvalue() == f"Creating with username {custom_username} and password Asd123!.\n"
+        f"Creating categories\n"
+        f"Creating ingredients\n"
+    )
+    assert User.objects.filter(username=custom_username).count() == 1
+    assert Category.objects.all().count() == 2
+    assert Ingredient.objects.all().count() == 4
+
+
+@pytest.mark.django_db
+def test_should_throw_exception_if_super_user_exists():
+    custom_username = "xpto"
+    get_user_model().objects.create_superuser(custom_username, None, "fake-password")
+
+    with pytest.raises(CommandError) as e:
+        call_command("seed_db", "--create-super-user", f"-u {custom_username}")
+
+    assert e.value.args[0] == "Super user xpto already exists"