"""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."))