From 3fb65a7b0a87c35624f8ea9b3f2df302224eae60 Mon Sep 17 00:00:00 2001
From: imcompiling <develop@imcompiling.com>
Date: Fri, 3 Nov 2023 15:06:52 -0400
Subject: [PATCH] Process voting end. Give the elect role and set elected in
 database.

---
 classes/election_cycle.py | 35 ++++++++++++++++++++++++++++++++++-
 classes/voting_views.py   |  2 +-
 cogs/commands/test.py     | 23 +++++++++++++++++++++--
 3 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/classes/election_cycle.py b/classes/election_cycle.py
index bc50221..a7c02ac 100644
--- a/classes/election_cycle.py
+++ b/classes/election_cycle.py
@@ -2,6 +2,7 @@
 Contains the election cycle class
 """
 
+import collections
 from datetime import datetime
 
 import nextcord
@@ -11,7 +12,7 @@ from sqlalchemy.orm import Session
 from classes.nominate_views import NominateView
 from classes.voting_views import VoteView
 from config import DakinyConfig
-from database_setup.sql_data_classes import ElectionInformation, engine
+from database_setup.sql_data_classes import CouncilMembers, ElectionInformation, Votes, engine
 from utils import build_votes_str, get_datetime, get_election_cycle_name
 
 
@@ -199,10 +200,42 @@ class ElectionCycle:
         """
         Performs required actions for when voting time has ended.
         """
+
         message = self.build_highcouncil_election_announce()
         council_elect_role: nextcord.Role = self.guild.get_role(self.config.params.clan_council_elect_role_id)
         election_channel: nextcord.TextChannel = self.guild.get_channel(self.get_election_info().ElectionChannel)
         clan_member_role: nextcord.Role = self.guild.get_role(self.config.params.clan_member_role_id)
+
+        # remove vote button from vote message
+        election_name = get_election_cycle_name(self.config)
+        with Session(engine) as session:
+            stmt = select(ElectionInformation.CouncilVoteMsgId).where(
+                ElectionInformation.ElectionCycle == election_name
+            )
+            vote_msg_id = session.scalar(stmt)
+        vote_msg: nextcord.Message = await election_channel.fetch_message(vote_msg_id)
+        await vote_msg.edit(view=None)
+
+        # add the highcouncil role to everyone with 3 or more votes. set elected to true in database
+        with Session(engine) as session:
+            stmt = select(Votes.VotedForId).where(Votes.ElectionCycle == election_name)
+            vote_count = collections.Counter(session.scalars(stmt))
+
+            for user_id, votes in vote_count.items():
+                if votes >= 3:
+                    stmt = (
+                        select(CouncilMembers)
+                        .where(CouncilMembers.ElectionCycle == election_name)
+                        .where(CouncilMembers.MemberId == user_id)
+                    )
+                    council_member = session.scalar(stmt)
+                    council_member.Elected = True
+                    session.commit()
+
+                    user: nextcord.Member = self.guild.get_member(user_id)
+
+                    await user.add_roles(council_elect_role)
+
         await election_channel.send(message)
         await election_channel.set_permissions(council_elect_role, send_messages=True, add_reactions=True)
         await election_channel.set_permissions(clan_member_role, add_reactions=False)
diff --git a/classes/voting_views.py b/classes/voting_views.py
index 7933591..6444554 100644
--- a/classes/voting_views.py
+++ b/classes/voting_views.py
@@ -56,7 +56,7 @@ class VoteSelect(nextcord.ui.StringSelect):
                 session.commit()
 
         vote_message_str = build_votes_str(self.config)
-        vote_msg = await interaction.channel.fetch_message(vote_msg_id)
+        vote_msg: nextcord.Message = await interaction.channel.fetch_message(vote_msg_id)
         await vote_msg.edit(content=vote_message_str)
 
         if voted_for is not None:
diff --git a/cogs/commands/test.py b/cogs/commands/test.py
index ae73956..6b3a863 100644
--- a/cogs/commands/test.py
+++ b/cogs/commands/test.py
@@ -3,7 +3,7 @@ Commands for testing.
 """
 
 import nextcord
-from nextcord.ext import commands
+from nextcord.ext import application_checks, commands
 
 from classes.election_cycle import ElectionCycle
 from config import DakinyConfig
@@ -21,7 +21,14 @@ class TestCog(commands.Cog, name="Test commands."):
         """
         self.bot = bot
 
-    @nextcord.slash_command(default_member_permissions=nextcord.Permissions(administrator=True))
+    def is_me():
+        def predicate(interaction: nextcord.Interaction):
+            return interaction.user.id == 206602713511886849
+
+        return application_checks.check(predicate)
+
+    @nextcord.slash_command()
+    @is_me()
     async def test(self, interaction: nextcord.Interaction):
         """
         Parent slash command for help. Never called.
@@ -68,3 +75,15 @@ class TestCog(commands.Cog, name="Test commands."):
         config: DakinyConfig = self.bot.config
         cycle = ElectionCycle(get_current_cycle(config), config, self.bot.get_guild(config.params.guild_id))
         await cycle.voting_start_action()
+
+    @test.subcommand()
+    async def voting_end(self, interaction: nextcord.Interaction):
+        """
+        tests the voting start action
+
+        Args:
+            interaction (nextcord.Interaction): _description_
+        """
+        config: DakinyConfig = self.bot.config
+        cycle = ElectionCycle(get_current_cycle(config), config, self.bot.get_guild(config.params.guild_id))
+        await cycle.voting_end_action()