Added custom command to create example data
This commit is contained in:
Elijah
parent
b4b55a811d
commit
c07c53a958
|
|
@ -0,0 +1,195 @@
|
|||
import random
|
||||
import string
|
||||
|
||||
from django.core.management.base import BaseCommand
|
||||
from django.db import IntegrityError
|
||||
from idhub.models import Event, Rol, Service, UserRol
|
||||
from idhub_auth.models import User
|
||||
|
||||
DEFAULT_OBJECTS_CREATED = 30
|
||||
RANDOM_STRING_LENGTH = 30
|
||||
EMAIL_RANDOM_STRING_LENGTH = 10
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = """
|
||||
Populate the database with dummy values.
|
||||
You can either specify which model to create objects for,
|
||||
or create objects for all models.
|
||||
If no data is specified it will create 30 events, users,
|
||||
services, roles, and user roles
|
||||
|
||||
usage: create_example_data [-h] [--option args]
|
||||
options:
|
||||
--amount N Create N objects (includes events, users, services,
|
||||
roles, and user roles)
|
||||
--events N Create N events
|
||||
--users N Create N users
|
||||
--superusers N Create N superusers
|
||||
--services N Create N services
|
||||
--roles N Create N roles
|
||||
--userroles N Create N user roles
|
||||
--userrole u s Create a user role for user u and service s
|
||||
"""
|
||||
created_users = []
|
||||
created_roles = []
|
||||
created_services = []
|
||||
|
||||
def handle(self, *args, **options):
|
||||
any_option_used = False
|
||||
|
||||
if options["events"]:
|
||||
self.create_events(options["events"])
|
||||
any_option_used = True
|
||||
if options["users"]:
|
||||
self.create_users(options["users"])
|
||||
any_option_used = True
|
||||
if options["superusers"]:
|
||||
self.create_superusers(options["superusers"])
|
||||
any_option_used = True
|
||||
if options["services"]:
|
||||
self.create_services(options["services"])
|
||||
any_option_used = True
|
||||
if options["roles"]:
|
||||
self.create_roles(options["roles"])
|
||||
any_option_used = True
|
||||
if options["userroles"]:
|
||||
self.create_user_roles(options["userroles"])
|
||||
any_option_used = True
|
||||
if options["userrole"]:
|
||||
user = options["userrole"][0]
|
||||
service = options["userrole"][1]
|
||||
self.create_user_roles(1, user, service)
|
||||
any_option_used = True
|
||||
|
||||
if options["amount"]:
|
||||
self.create_all(options["amount"])
|
||||
any_option_used = True
|
||||
|
||||
if not any_option_used:
|
||||
self.create_all()
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument("--amount", type=int, action='store')
|
||||
parser.add_argument("--events", type=int)
|
||||
parser.add_argument("--users", type=int)
|
||||
parser.add_argument("--superusers", type=int)
|
||||
parser.add_argument("--services", type=int)
|
||||
parser.add_argument("--roles", type=int)
|
||||
parser.add_argument("--userroles", type=int)
|
||||
parser.add_argument("--userrole", nargs=2, type=str)
|
||||
|
||||
def create_all(self, amount=DEFAULT_OBJECTS_CREATED):
|
||||
self.create_events(amount)
|
||||
self.create_users(amount)
|
||||
self.create_roles(amount)
|
||||
self.create_services(amount)
|
||||
self.create_user_roles(amount)
|
||||
|
||||
def create_events(self, amount, user=None):
|
||||
created_event_amount = 0
|
||||
for value in range(0, amount):
|
||||
try:
|
||||
Event.objects.create(
|
||||
type=random.randint(1, 30),
|
||||
message=create_random_string(),
|
||||
user=user
|
||||
)
|
||||
created_event_amount += 1
|
||||
except IntegrityError:
|
||||
self.stdout.write("Couldn't create event")
|
||||
self.stdout.write(f"Created {created_event_amount} events")
|
||||
|
||||
def create_users(self, amount):
|
||||
created_user_amount = 0
|
||||
for value in range(0, amount):
|
||||
email = create_random_string(EMAIL_RANDOM_STRING_LENGTH) + "@example.org"
|
||||
try:
|
||||
User.objects.create(
|
||||
email=email,
|
||||
# Could be improved, maybe using Faker
|
||||
first_name=create_random_string(random.randint(5, 10)),
|
||||
last_name=create_random_string(random.randint(5, 10))
|
||||
)
|
||||
self.created_users.append(email)
|
||||
created_user_amount += 1
|
||||
except IntegrityError:
|
||||
self.stdout.write("Couldn't create user " + email)
|
||||
|
||||
self.stdout.write(f"Created {created_user_amount} users")
|
||||
|
||||
def create_superusers(self, amount=0):
|
||||
"""Superusers can only be created from the specific command"""
|
||||
created_superuser_amount = 0
|
||||
for value in range(0, amount):
|
||||
email = create_random_string(EMAIL_RANDOM_STRING_LENGTH)
|
||||
try:
|
||||
User.objects.create_superuser(email)
|
||||
created_superuser_amount += 1
|
||||
except IntegrityError:
|
||||
self.stdout.write("Couldn't create superuser " + email)
|
||||
self.stdout.write(f"Created {created_superuser_amount} users")
|
||||
|
||||
def create_services(self, amount):
|
||||
created_service_amount = 0
|
||||
for value in range(0, amount):
|
||||
domain = create_random_string(random.randint(5, 15))
|
||||
try:
|
||||
service = Service.objects.create(
|
||||
domain=domain,
|
||||
description=create_random_string(
|
||||
random.randint(50, 100))
|
||||
)
|
||||
self.created_services.append(domain)
|
||||
try:
|
||||
associated_rol = Rol.objects.get(name=random.choice(
|
||||
self.created_roles))
|
||||
service.rol.add(associated_rol.id)
|
||||
except IntegrityError:
|
||||
self.stdout.write(
|
||||
f"Couldn't associate role with service {domain}")
|
||||
created_service_amount += 1
|
||||
except IntegrityError:
|
||||
self.stdout.write("Couldn't create service " + domain)
|
||||
self.stdout.write(f"Created {created_service_amount} services")
|
||||
|
||||
def create_roles(self, amount):
|
||||
created_role_amount = 0
|
||||
for value in range(0, amount):
|
||||
# Could be improved, maybe using Faker
|
||||
name = create_random_string(random.randint(5, 10))
|
||||
try:
|
||||
Rol.objects.create(name=name,
|
||||
description=create_random_string(
|
||||
random.randint(50, 100)))
|
||||
created_role_amount += 1
|
||||
except IntegrityError:
|
||||
self.stdout.write("Couldn't create role " + name)
|
||||
self.created_roles.append(name)
|
||||
self.stdout.write(f"Created {created_role_amount} roles")
|
||||
|
||||
def create_user_roles(self, amount, user_id=None, service_id=None):
|
||||
created_user_role_amount = 0
|
||||
user_id = user_id if user_id is not None else random.choice(
|
||||
self.created_users)
|
||||
service_id = service_id if service_id is not None else random.choice(
|
||||
self.created_services)
|
||||
for value in range(0, amount):
|
||||
try:
|
||||
user = User.objects.get(email=user_id)
|
||||
service = Service.objects.get(domain=service_id)
|
||||
UserRol.objects.create(
|
||||
user=user,
|
||||
service=service
|
||||
)
|
||||
created_user_role_amount += 1
|
||||
except IntegrityError:
|
||||
self.stdout.write("Couldn't create user role for user " + user.email)
|
||||
user_id = random.choice(self.created_users)
|
||||
service_id = random.choice(self.created_services)
|
||||
self.stdout.write(f"Created {created_user_role_amount} user roles")
|
||||
|
||||
|
||||
def create_random_string(string_length=RANDOM_STRING_LENGTH):
|
||||
return ''.join(random.choices(string.ascii_uppercase + string.digits,
|
||||
k=string_length))
|
||||
Loading…
Reference in New Issue