From e851adee01e9cb2fae147c3fb16ca2161b753d74 Mon Sep 17 00:00:00 2001
From: Connor Nelson <Connor@ConnorNelson.com>
Date: Thu, 30 Nov 2023 11:03:08 -0700
Subject: [PATCH] Refactor: Add get_user_belts util

---
 dojo_plugin/__init__.py            | 17 ++---------------
 dojo_plugin/models/__init__.py     |  2 +-
 dojo_plugin/utils/dojo/__init__.py | 16 ++++++++++++++++
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/dojo_plugin/__init__.py b/dojo_plugin/__init__.py
index 0a9fb51e8..347010572 100644
--- a/dojo_plugin/__init__.py
+++ b/dojo_plugin/__init__.py
@@ -19,6 +19,7 @@
 from .config import DOJO_HOST, bootstrap
 from .utils import unserialize_user_flag, render_markdown
 from .utils.discord import get_discord_user, get_discord_roles, add_role, send_message
+from .utils.dojo import get_user_belts
 from .pages.dojos import dojos, dojos_override
 from .pages.dojo import dojo
 from .pages.workspace import workspace
@@ -45,24 +46,10 @@ def solve(cls, user, team, challenge, request):
         if not discord_user:
             return
 
-        # TODO: Get this from dojo.yml
-        belts = {
-            "CSE 365 - Spring 2023": "Orange Belt",
-            "CSE 466 - Fall 2022": "Yellow Belt",
-            "CSE 494 - Spring 2023": "Blue Belt",
-        }
         discord_roles = get_discord_roles()
-
-        dojos = Dojos.query.join(DojoChallenges).filter(Dojos.official, DojoChallenges.challenge == challenge)
-        for dojo in dojos:
-            belt = belts.get(dojo.name)
-            if not belt:
-                continue
-            if not dojo.completed(user):
-                continue
+        for belt in get_user_belts(user):
             if discord_roles.get(belt) in discord_user["roles"]:
                 continue
-
             user_mention = f"<@{discord_user['user']['id']}>"
             message = f"{user_mention} earned their {belt}! :tada:"
             print(message, flush=True)
diff --git a/dojo_plugin/models/__init__.py b/dojo_plugin/models/__init__.py
index e33aaa152..215be4597 100644
--- a/dojo_plugin/models/__init__.py
+++ b/dojo_plugin/models/__init__.py
@@ -195,7 +195,7 @@ def solves(self, **kwargs):
         return DojoChallenges.solves(dojo=self, **kwargs)
 
     def completed(self, user):
-        return self.solves(user=user).count() == len(self.challenges)
+        return self.solves(user=user, ignore_visibility=True, ignore_admins=False).count() == len(self.challenges)
 
     def is_admin(self, user=None):
         if user is None:
diff --git a/dojo_plugin/utils/dojo/__init__.py b/dojo_plugin/utils/dojo/__init__.py
index 0f95a1044..f74772c7c 100644
--- a/dojo_plugin/utils/dojo/__init__.py
+++ b/dojo_plugin/utils/dojo/__init__.py
@@ -397,3 +397,19 @@ def get_current_dojo_challenge(user=None):
                 DojoChallenges.dojo == Dojos.from_id(container.labels.get("dojo.dojo_id")).first())
         .first()
     )
+
+
+def get_user_belts(user):
+    # TODO: Get this from dojo.yml
+    belts = {
+        "Orange Belt": ["cse365-s2023", "cse365-f2023"],
+        "Yellow Belt": ["cse466-f2022", "cse466-f2023"],
+        "Blue Belt": ["cse494-s2023"],
+    }
+    result = []
+    for belt, dojo_ids in belts.items():
+        dojos = Dojos.query.filter(Dojos.official, Dojos.id.in_(dojo_ids))
+        if not any(dojo.completed(user) for dojo in dojos):
+            break
+        result.append(belt)
+    return result
\ No newline at end of file