65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
"""Command to create test data fixture for teams."""
|
|
|
|
import random
|
|
import json
|
|
from string import ascii_uppercase
|
|
|
|
from django.core.management.base import BaseCommand
|
|
from django.db.utils import IntegrityError
|
|
from mainapp.models import Team, BLOCKED_SECTION_LETTERS
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Creates a fixture file for Team objects"
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument("amount_of_teams", type=int, help="Number of teams to create")
|
|
|
|
def handle(self, *args, **options):
|
|
|
|
existing_teams = Team.objects.all()
|
|
|
|
# Available sections
|
|
available_sections = [
|
|
char for char in [*ascii_uppercase]
|
|
if char.lower() not in BLOCKED_SECTION_LETTERS
|
|
and char not in [team.section_letter for team in existing_teams]
|
|
]
|
|
|
|
if not available_sections:
|
|
self.stdout.write(self.style.ERROR(
|
|
f"There are no available sections for new teams."
|
|
))
|
|
return
|
|
|
|
teams_amount = options["amount_of_teams"]
|
|
max_teams = len(available_sections)
|
|
|
|
if teams_amount > max_teams:
|
|
self.stdout.write(self.style.ERROR(
|
|
f"Number of teams is too large [{teams_amount}/{max_teams}]."
|
|
))
|
|
return
|
|
|
|
# Create the new teams (this will create them in the database)#
|
|
new_teams = []
|
|
for i in range(teams_amount):
|
|
section_letter = random.choice(available_sections)
|
|
available_sections.remove(section_letter)
|
|
team = Team.objects.create(section_letter=section_letter)
|
|
new_teams.append(team)
|
|
|
|
teams_fixture = [{
|
|
"model": "mainapp.team",
|
|
"pk": team.pk,
|
|
"fields": {"section_letter": team.section_letter}
|
|
} for team in new_teams]
|
|
|
|
# Remove the teams from the database
|
|
for team in new_teams:
|
|
team.delete()
|
|
|
|
with open("src/mainapp/fixtures/teams_fixture.json", "w") as file:
|
|
file.write(json.dumps(teams_fixture, indent=4))
|
|
self.stdout.write(self.style.SUCCESS("Created teams_fixture.json."))
|