Refactor badges update job: modularize handlers, introduce base class, and add tests

backend/apps/nest/Makefile

@@ -1,3 +1,3 @@
 nest-update-badges:
 	@echo "Updating Nest user badges"
-	@CMD="python manage.py nest_update_badges" $(MAKE) exec-backend-command
+	@CMD="python manage.py nest_update_badges $(ARGS)" $(MAKE) exec-backend-command

backend/apps/nest/badges/__init__.py

@@ -0,0 +1,9 @@
+"""Badge handlers package for managing user badges."""
+
+from apps.nest.badges.project_leader_badge import OWASPProjectLeaderBadgeHandler
+from apps.nest.badges.staff_badge import OWASPStaffBadgeHandler
+
+__all__ = [
+    "OWASPProjectLeaderBadgeHandler",
+    "OWASPStaffBadgeHandler",
+]

backend/apps/nest/badges/base.py

@@ -0,0 +1,86 @@
+"""Base badge handler module."""
+
+import logging
+from abc import ABC, abstractmethod
+
+from django.db.models import QuerySet
+
+from apps.github.models.user import User
+from apps.nest.models.badge import Badge
+from apps.nest.models.user_badge import UserBadge
+
+logger = logging.getLogger(__name__)
+
+
+class BaseBadgeHandler(ABC):
+    """Base class for handling badge updates."""
+
+    name: str
+    description: str
+    css_class: str
+    weight: int
+
+    def __init__(self, stdout=None, style=None):
+        """Initialize the handler."""
+        self.stdout = stdout
+        self.style = style
+
+    @abstractmethod
+    def get_eligible_users(self) -> QuerySet[User]:
+        """Return a QuerySet of users who should currently have this badge."""
+
+    def get_badge_defaults(self) -> dict:
+        """Return the default attributes for the badge."""
+        return {
+            "description": self.description,
+            "css_class": self.css_class,
+            "weight": self.weight,
+        }
+
+    def _log(self, message, style_func=None):
+        """Log to both file logger and stdout if available."""
+        logger.info(message)
+        if self.stdout:
+            if style_func:
+                message = style_func(message)
+            self.stdout.write(message)
+
+    def process(self):
+        """Execute the badge sync process."""
+        if not self.name:
+            msg = "Badge name must be defined."
+            raise ValueError(msg)
+
+        badge, created = Badge.objects.get_or_create(
+            name=self.name,
+            defaults=self.get_badge_defaults(),
+        )
+
+        if created:
+            self._log(f"Created badge: '{badge.name}'")
+
+        user_without_badge = self.get_eligible_users()
+        users_to_add = user_without_badge.exclude(
+            user_badges__badge=badge,
+            user_badges__is_active=True,
+        )
+
+        added_count = 0
+        for user in users_to_add:
+            user_badge, created = UserBadge.objects.get_or_create(user=user, badge=badge)
+            if not user_badge.is_active:
+                user_badge.is_active = True
+                user_badge.save(update_fields=["is_active"])
+                added_count += 1
+
+        self._log(f"Added '{self.name}' badge to {added_count} users")
+
+        users_to_remove_badge = UserBadge.objects.filter(
+            badge=badge,
+            is_active=True,
+        ).exclude(user__in=user_without_badge)
+
+        count = users_to_remove_badge.count()
+        if count:
+            users_to_remove_badge.update(is_active=False)
+            self._log(f"Removed '{self.name}' badge from {count} users")

backend/apps/nest/badges/project_leader_badge.py

@@ -0,0 +1,34 @@
+"""Handler for the OWASP Project Leader badge."""
+
+from django.contrib.contenttypes.models import ContentType
+from django.db.models import QuerySet
+
+from apps.github.models.user import User
+from apps.nest.badges.base import BaseBadgeHandler
+from apps.owasp.models.entity_member import EntityMember
+from apps.owasp.models.project import Project
+
+
+class OWASPProjectLeaderBadgeHandler(BaseBadgeHandler):
+    """Handler for managing the OWASP Project Leader badge."""
+
+    name = "OWASP Project Leader"
+    description = "Official OWASP Project Leader"
+    css_class = "fa-user-shield"
+    weight = 90
+
+    def get_eligible_users(self) -> QuerySet[User]:
+        """Get all users who should have the Project Leader badge.
+
+        Returns:
+            QuerySet of users who are project leaders.
+
+        """
+        leader_ids = EntityMember.objects.filter(
+            entity_type=ContentType.objects.get_for_model(Project),
+            role=EntityMember.Role.LEADER,
+            is_active=True,
+            is_reviewed=True,
+            member__isnull=False,
+        ).values_list("member_id", flat=True)
+        return User.objects.filter(id__in=leader_ids).distinct()

backend/apps/nest/badges/staff_badge.py

@@ -0,0 +1,24 @@
+"""Handler for the OWASP Staff badge."""
+
+from django.db.models import QuerySet
+
+from apps.github.models.user import User
+from apps.nest.badges.base import BaseBadgeHandler
+
+
+class OWASPStaffBadgeHandler(BaseBadgeHandler):
+    """Handler for managing the OWASP Staff badge."""
+
+    name = "OWASP Staff"
+    description = "Official OWASP Staff"
+    css_class = "fa-user-shield"
+    weight = 100
+
+    def get_eligible_users(self) -> QuerySet[User]:
+        """Get all users who should have the OWASP Staff badge.
+
+        Returns:
+            QuerySet of users with is_owasp_staff=True.
+
+        """
+        return User.objects.filter(is_owasp_staff=True)

backend/apps/nest/management/commands/nest_update_badges.py

@@ -2,152 +2,51 @@
 
 import logging
 
-from django.contrib.contenttypes.models import ContentType
 from django.core.management.base import BaseCommand
 
-from apps.github.models.user import User
-from apps.nest.models.badge import Badge
-from apps.nest.models.user_badge import UserBadge
-from apps.owasp.models.entity_member import EntityMember
-from apps.owasp.models.project import Project
+from apps.nest.badges import OWASPProjectLeaderBadgeHandler, OWASPStaffBadgeHandler
 
 logger = logging.getLogger(__name__)
 
-OWASP_STAFF_BADGE_NAME = "OWASP Staff"
-OWASP_PROJECT_LEADER_BADGE_NAME = "OWASP Project Leader"
-
 
 class Command(BaseCommand):
     """Sync badges for users based on their roles and attributes."""
 
     help = "Sync badges for users based on their roles and attributes"
 
-    def handle(self, *args, **options):
-        """Execute the command."""
-        self.stdout.write("Syncing user badges...")
-        self.update_owasp_staff_badge()
-        self.update_owasp_project_leader_badge()
-        self.stdout.write(self.style.SUCCESS("User badges sync completed"))
-
-    def update_owasp_staff_badge(self):
-        """Sync OWASP Staff badge for users."""
-        # Get or create the OWASP Staff badge
-        badge, created = Badge.objects.get_or_create(
-            name=OWASP_STAFF_BADGE_NAME,
-            defaults={
-                "description": "Official OWASP Staff",
-                "css_class": "fa-user-shield",
-                "weight": 100,  # High weight for importance
-            },
-        )
-
-        if created:
-            logger.info("Created '%s' badge", OWASP_STAFF_BADGE_NAME)
-            self.stdout.write(f"Created badge: {badge.name}")
-
-        # Assign badge to employees who don't have it.
-        employees_without_badge = User.objects.filter(
-            is_owasp_staff=True,
-        ).exclude(
-            user_badges__badge=badge,
-        )
-        count = employees_without_badge.count()
-
-        if count:
-            for user in employees_without_badge:
-                user_badge, created = UserBadge.objects.get_or_create(user=user, badge=badge)
-                if not user_badge.is_active:
-                    user_badge.is_active = True
-                    user_badge.save(update_fields=["is_active"])
-
-        logger.info("Added '%s' badge to %s users", OWASP_STAFF_BADGE_NAME, count)
-        self.stdout.write(f"Added badge to {count} employees")
-
-        # Remove badge from non-OWASP employees.
-        non_employees = User.objects.filter(
-            is_owasp_staff=False,
-            user_badges__badge=badge,
-        ).distinct()
-        removed_count = non_employees.count()
-
-        if removed_count:
-            UserBadge.objects.filter(
-                user_id__in=non_employees.values_list("id", flat=True),
-                badge=badge,
-            ).update(is_active=False)
-
-        logger.info("Removed '%s' badge from %s users", OWASP_STAFF_BADGE_NAME, removed_count)
-        self.stdout.write(f"Removed badge from {removed_count} non-employees")
-
-    def update_owasp_project_leader_badge(self):
-        """Sync OWASP Project Leader badge for users."""
-        badge, created = Badge.objects.get_or_create(
-            name=OWASP_PROJECT_LEADER_BADGE_NAME,
-            defaults={
-                "description": "Official OWASP Project Leader",
-                "css_class": "fa-user-shield",
-                "weight": 90,
-            },
-        )
-
-        if created:
-            logger.info("Created '%s' badge", OWASP_PROJECT_LEADER_BADGE_NAME)
-            self.stdout.write(f"Created badge: {badge.name}")
-
-        project_leaders_without_badge = (
-            User.objects.filter(
-                id__in=EntityMember.objects.filter(
-                    entity_type=ContentType.objects.get_for_model(Project),
-                    role=EntityMember.Role.LEADER,
-                    is_active=True,
-                    is_reviewed=True,
-                    member__isnull=False,
-                ).values_list("member_id", flat=True),
-            )
-            .distinct()
-            .exclude(
-                user_badges__badge=badge,
-            )
-        )
-
-        count = project_leaders_without_badge.count()
-
-        if count:
-            for user in project_leaders_without_badge:
-                user_badge, created = UserBadge.objects.get_or_create(
-                    user=user,
-                    badge=badge,
-                )
-                if not user_badge.is_active:
-                    user_badge.is_active = True
-                    user_badge.save(update_fields=["is_active"])
-
-        logger.info("Added '%s' badge to %s users", OWASP_PROJECT_LEADER_BADGE_NAME, count)
-        self.stdout.write(f"Added badge to {count} project leaders")
-
-        # Remove badge from users who are no longer project leaders.
-        current_leader_ids = (
-            EntityMember.objects.filter(
-                entity_type=ContentType.objects.get_for_model(Project),
-                role=EntityMember.Role.LEADER,
-                is_active=True,
-                is_reviewed=True,
-                member__isnull=False,
-            )
-            .values_list("member_id", flat=True)
-            .distinct()
+    BADGE_HANDLERS = {
+        "staff": OWASPStaffBadgeHandler,
+        "leader": OWASPProjectLeaderBadgeHandler,
+    }
+
+    def add_arguments(self, parser):
+        """Add command arguments."""
+        parser.add_argument(
+            "--handler",
+            type=str,
+            choices=self.BADGE_HANDLERS.keys(),
+            help="Specific badge handler to run. Runs all if omitted.",
+            required=False,
         )
 
-        non_leaders = UserBadge.objects.filter(
-            badge=badge,
-            is_active=True,
-        ).exclude(user_id__in=current_leader_ids)
-        removed_count = non_leaders.count()
-
-        if removed_count:
-            non_leaders.update(is_active=False)
+    def handle(self, *args, **options):
+        """Execute the command."""
+        handler_key = options.get("handler")
+
+        if handler_key:
+            self.stdout.write(f"Syncing specific badge: {handler_key}...")
+            handlers_to_run = [self.BADGE_HANDLERS[handler_key]]
+        else:
+            self.stdout.write("Syncing ALL user badges...")
+            handlers_to_run = self.BADGE_HANDLERS.values()
+
+        for handler_class in handlers_to_run:
+            try:
+                self.stdout.write(f"Processing badge: {handler_class.name}")
+                handler = handler_class(stdout=self.stdout, style=self.style)
+                handler.process()
+            except Exception:
+                logger.exception("Error processing badge %s", handler_class.name)
+                self.stdout.write(self.style.ERROR(f"Failed to update {handler_class.name}"))
 
-        logger.info(
-            "Removed '%s' badge from %s users", OWASP_PROJECT_LEADER_BADGE_NAME, removed_count
-        )
-        self.stdout.write(f"Removed badge from {removed_count} non-project leaders")
+        self.stdout.write(self.style.SUCCESS("User badges sync completed"))