38686-vm/core/management/commands/setup_test_data.py

# === SETUP TEST DATA MANAGEMENT COMMAND ===
# Creates sample workers, projects, teams, and work logs for testing.
# Run this once after deploying: python manage.py setup_test_data
#
# This is useful when the Django admin panel isn't accessible (e.g. on
# Flatlogic's Cloud Run deployment where admin static files may not load).

import datetime
from decimal import Decimal

from django.core.management.base import BaseCommand
from django.contrib.auth.models import User
from django.utils import timezone

from core.models import Project, Worker, Team, WorkLog


class Command(BaseCommand):
    help = 'Creates sample workers, projects, teams, and work logs for testing'

    def handle(self, *args, **options):
        # --- Create an admin superuser (if not exists) ---
        admin_user, created = User.objects.get_or_create(
            username='admin',
            defaults={
                'is_staff': True,
                'is_superuser': True,
                'first_name': 'Admin',
                'email': 'admin@foxfitt.co.za',
            }
        )
        if created:
            admin_user.set_password('admin123')
            admin_user.save()
            self.stdout.write(self.style.SUCCESS('Created admin user (password: admin123)'))
        else:
            # Make sure existing admin has staff/superuser flags
            if not admin_user.is_staff or not admin_user.is_superuser:
                admin_user.is_staff = True
                admin_user.is_superuser = True
                admin_user.save()
                self.stdout.write('Updated admin user to have staff + superuser flags')
            else:
                self.stdout.write('Admin user already exists')

        # --- Create a supervisor user ---
        supervisor, created = User.objects.get_or_create(
            username='supervisor1',
            defaults={
                'is_staff': False,
                'first_name': 'John',
                'last_name': 'Supervisor',
            }
        )
        if created:
            supervisor.set_password('super123')
            supervisor.save()
            self.stdout.write(self.style.SUCCESS('Created supervisor user (password: super123)'))
        else:
            self.stdout.write('Supervisor user already exists')

        # --- Create Projects ---
        project_names = [
            'Kalkbult Solar Farm',
            'De Aar Wind Farm',
            'Prieska Solar Plant',
        ]
        projects = []
        for name in project_names:
            proj, created = Project.objects.get_or_create(
                name=name,
                defaults={'active': True}
            )
            # Assign supervisor to the project
            proj.supervisors.add(supervisor)
            projects.append(proj)
            if created:
                self.stdout.write(self.style.SUCCESS(f'  Created project: {name}'))
            else:
                self.stdout.write(f'  Project already exists: {name}')

        # --- Create Workers ---
        worker_data = [
            {'name': 'Thabo Mokoena', 'id_number': '9001015000080', 'salary': Decimal('8000.00')},
            {'name': 'Sipho Ndlovu', 'id_number': '8805125000081', 'salary': Decimal('7500.00')},
            {'name': 'Lerato Dlamini', 'id_number': '9203220000082', 'salary': Decimal('7000.00')},
            {'name': 'Bongani Zulu', 'id_number': '8510305000083', 'salary': Decimal('8500.00')},
            {'name': 'Nomsa Khumalo', 'id_number': '9106185000084', 'salary': Decimal('7200.00')},
            {'name': 'David Botha', 'id_number': '8707125000085', 'salary': Decimal('9000.00')},
        ]
        workers = []
        for wd in worker_data:
            worker, created = Worker.objects.get_or_create(
                name=wd['name'],
                defaults={
                    'id_number': wd['id_number'],
                    'monthly_salary': wd['salary'],
                    'active': True,
                    'employment_date': datetime.date(2024, 1, 15),
                }
            )
            workers.append(worker)
            if created:
                self.stdout.write(self.style.SUCCESS(f'  Created worker: {wd["name"]} (R{wd["salary"]}/month)'))
            else:
                self.stdout.write(f'  Worker already exists: {wd["name"]}')

        # --- Create Teams ---
        team_a, created = Team.objects.get_or_create(
            name='Team Alpha',
            defaults={'supervisor': supervisor, 'active': True}
        )
        if created:
            team_a.workers.set(workers[:3])  # First 3 workers
            self.stdout.write(self.style.SUCCESS('  Created Team Alpha (3 workers)'))
        else:
            self.stdout.write('  Team Alpha already exists')

        team_b, created = Team.objects.get_or_create(
            name='Team Bravo',
            defaults={'supervisor': supervisor, 'active': True}
        )
        if created:
            team_b.workers.set(workers[3:])  # Last 3 workers
            self.stdout.write(self.style.SUCCESS('  Created Team Bravo (3 workers)'))
        else:
            self.stdout.write('  Team Bravo already exists')

        # --- Create Work Logs (last 2 weeks) ---
        today = timezone.now().date()
        logs_created = 0

        for days_ago in range(14, 0, -1):
            log_date = today - datetime.timedelta(days=days_ago)

            # Skip weekends
            if log_date.weekday() >= 5:
                continue

            # Alternate between projects
            project = projects[days_ago % len(projects)]

            # Create a work log with some workers
            log_workers = workers[:4] if days_ago % 2 == 0 else workers[2:]

            # Check if this log already exists
            existing = WorkLog.objects.filter(date=log_date, project=project).first()
            if existing:
                continue

            # Set overtime on some days
            ot = Decimal('0.00')
            if days_ago % 3 == 0:
                ot = Decimal('0.50')  # Half day overtime every 3rd day
            elif days_ago % 5 == 0:
                ot = Decimal('0.25')  # Quarter day overtime every 5th day

            worklog = WorkLog.objects.create(
                date=log_date,
                project=project,
                team=team_a if days_ago % 2 == 0 else team_b,
                supervisor=supervisor,
                overtime_amount=ot,
            )
            worklog.workers.set(log_workers)
            logs_created += 1

        self.stdout.write(self.style.SUCCESS(f'  Created {logs_created} work logs'))

        # --- Summary ---
        self.stdout.write('')
        self.stdout.write(self.style.SUCCESS('=== Test data setup complete! ==='))
        self.stdout.write(f'  Admin login:      admin / admin123')
        self.stdout.write(f'  Supervisor login:  supervisor1 / super123')
        self.stdout.write(f'  Projects: {Project.objects.filter(active=True).count()}')
        self.stdout.write(f'  Workers:  {Worker.objects.filter(active=True).count()}')
        self.stdout.write(f'  Teams:    {Team.objects.filter(active=True).count()}')
        self.stdout.write(f'  WorkLogs: {WorkLog.objects.count()}')
        self.stdout.write('')
        self.stdout.write('  Now log in as "admin" and go to /payroll/ to test the dashboard!')