Add management command to populate test data for OctoFit collections

Author: AlgorithmicThoughts

octofit-tracker/backend/octofit_tracker/management/commands/populate_db.py

@@ -0,0 +1,50 @@
+from django.core.management.base import BaseCommand
+from octofit_tracker.models import User, Team, Activity, Leaderboard, Workout
+from djongo import models
+
+class Command(BaseCommand):
+    help = 'Populate the octofit_db database with test data'
+
+    def handle(self, *args, **options):
+        # Delete existing data
+        Activity.objects.all().delete()
+        Workout.objects.all().delete()
+        Leaderboard.objects.all().delete()
+        User.objects.all().delete()
+        Team.objects.all().delete()
+
+        # Create teams
+        marvel = Team.objects.create(name='Marvel')
+        dc = Team.objects.create(name='DC')
+
+        # Create users
+        users = [
+            User.objects.create(name='Spider-Man', email='spiderman@marvel.com', team=marvel),
+            User.objects.create(name='Iron Man', email='ironman@marvel.com', team=marvel),
+            User.objects.create(name='Wonder Woman', email='wonderwoman@dc.com', team=dc),
+            User.objects.create(name='Batman', email='batman@dc.com', team=dc),
+        ]
+
+        # Create activities
+        Activity.objects.create(user=users[0], type='Running', duration=30)
+        Activity.objects.create(user=users[1], type='Cycling', duration=45)
+        Activity.objects.create(user=users[2], type='Swimming', duration=60)
+        Activity.objects.create(user=users[3], type='Yoga', duration=20)
+
+        # Create workouts
+        workout1 = Workout.objects.create(name='Hero HIIT', description='High intensity interval training for heroes')
+        workout2 = Workout.objects.create(name='Power Yoga', description='Yoga for strength and flexibility')
+        workout1.suggested_for.set([users[0], users[1]])
+        workout2.suggested_for.set([users[2], users[3]])
+
+        # Create leaderboard
+        Leaderboard.objects.create(team=marvel, points=100)
+        Leaderboard.objects.create(team=dc, points=90)
+
+        # Ensure unique index on email
+        from pymongo import MongoClient
+        client = MongoClient('localhost', 27017)
+        db = client['octofit_db']
+        db.users.create_index('email', unique=True)
+
+        self.stdout.write(self.style.SUCCESS('Test data populated successfully!'))