Merge pull request #31 from EboMike/cleanup

Moved more common code into helpers.

handlers/game/game.py

@@ -1,4 +1,7 @@
 from sanic import Request
+
+
+from helpers.gamehelper import get_team_rosters, get_player_display_names, get_matchmaking_teams
 from shared import app
 from utils import render_template, is_admin
 from db.game import EntityEnds, EntityStarts
@@ -24,28 +27,14 @@ async def get_laserballstats(entity) -> Optional[LaserballStats]:
 @sentry_trace
 async def game_index(request: Request, type: str, id: int) -> str:
     if type == "sm5":
-        game: SM5Game = await SM5Game.filter(id=id).prefetch_related("entity_starts").first()
+        game: SM5Game = await SM5Game.filter(id=id).prefetch_related("entity_starts", "entity_ends").first()
 
         if not game:
             raise exceptions.NotFound("Not found: Invalid game ID")
 
-        players_matchmake_team1 = []
-        players_matchmake_team2 = []
-        entity_starts: List[EntityStarts] = game.entity_starts
-        for i, player in enumerate(entity_starts):
-            if player.type != "player":
-                continue
+        team_rosters = await get_team_rosters(game.entity_starts, game.entity_ends)
 
-            if (await player.team).enum == Team.RED:
-                if player.entity_id.startswith("@"):
-                    players_matchmake_team1.append(player.name)
-                else:
-                    players_matchmake_team1.append(player.entity_id)
-            elif (await player.team).enum in [Team.BLUE, Team.GREEN]:
-                if player.entity_id.startswith("@"):
-                    players_matchmake_team2.append(player.name)
-                else:
-                    players_matchmake_team2.append(player.entity_id)
+        players_matchmake_team1, players_matchmake_team2 = get_matchmaking_teams(team_rosters)
 
         return await render_template(
             request, "game/sm5.html",
@@ -63,29 +52,15 @@ async def game_index(request: Request, type: str, id: int) -> str:
             is_admin=is_admin(request)
         )
     elif type == "laserball":
-        game = await LaserballGame.filter(id=id).prefetch_related("entity_starts").first()
+        game = await LaserballGame.filter(id=id).prefetch_related("entity_starts", "entity_ends").first()
 
         if not game:
             raise exceptions.NotFound("Not found: Invalid game ID")
 
-        players_matchmake_team1 = []
-        players_matchmake_team2 = []
-        entity_starts: List[EntityStarts] = game.entity_starts
-        for i, player in enumerate(entity_starts):
-            if player.type != "player":
-                continue
+        team_rosters = await get_team_rosters(game.entity_starts, game.entity_ends)
+
+        players_matchmake_team1, players_matchmake_team2 = get_matchmaking_teams(team_rosters)
 
-            if (await player.team).enum == Team.RED:
-                if player.entity_id.startswith("@"):
-                    players_matchmake_team1.append(player.name)
-                else:
-                    players_matchmake_team1.append(player.entity_id)
-            elif (await player.team).enum in [Team.BLUE, Team.GREEN]:
-                if player.entity_id.startswith("@"):
-                    players_matchmake_team2.append(player.name)
-                else:
-                    players_matchmake_team2.append(player.entity_id)
-
         return await render_template(
             request, "game/laserball.html",
             game=game, get_entity_end=get_entity_end, get_laserballstats=get_laserballstats,

helpers/gamehelper.py

@@ -3,9 +3,14 @@
 This is different from statshelper in that it does not calculate stats, it
 only deals with getting and extracting and visualizing data.
 """
-from typing import List
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import List, Optional, Dict
 
-from db.types import PlayerStateDetailType
+from sanic import exceptions
+
+from db.game import EntityStarts, EntityEnds, Teams
+from db.types import PlayerStateDetailType, Team
 
 """Map of every possible player state and the display name for it in SM5 games.
 
@@ -34,8 +39,83 @@
 }
 
 
+@dataclass
+class PlayerInfo:
+    """Information about a player in one particular game."""
+    entity_start: EntityStarts
+    entity_end: Optional[EntityEnds]
+    display_name: str
+
+
 def get_players_from_team(all_players: List[dict], team_index: int) -> List[dict]:
-    """Returns subset of the list of players - only those in the given team."""
+    """Returns subset of the list of players - only those in the given team.
+
+    Each player object is a dict that has a "team" member with a team_index.
+    """
     return [
         player for player in all_players if player["team"] == team_index
     ]
+
+
+async def get_team_rosters(entity_starts: List[EntityStarts], entity_ends: List[EntityEnds]) -> dict[
+    Team, List[PlayerInfo]]:
+    """Returns a dict with each team and a list of players in each.
+
+    Non-player entities will be ignored. The values will be a list of names, either the
+    name if the player is not logged in, or the entity ID if the player is logged in.
+    """
+    result = defaultdict(list)
+
+    entity_ends_dict = {
+        (await entity.entity).id: entity for entity in entity_ends
+    }
+
+    for player in entity_starts:
+        if player.type != "player":
+            continue
+
+        player_team = (await player.team).enum
+        team_roster = result[player_team]
+
+        entity_end = entity_ends_dict[player.id] if player.id in entity_ends_dict else None
+        display_name = player.name if player.entity_id.startswith("@") else player.entity_id
+
+        team_roster.append(PlayerInfo(entity_start=player, entity_end=entity_end, display_name=display_name))
+
+    return result
+
+
+def get_player_display_names(players: List[PlayerInfo]) -> List[str]:
+    """Extracts all the display names from a list of players."""
+    return [player.display_name for player in players]
+
+
+def get_matchmaking_teams(team_rosters: Dict[Team, List[PlayerInfo]]) -> (
+    List[str], List[str]):
+    """Returns display names for each player in both teams.
+
+    The first returned list will always be the red team unless there is no red
+    team. The second one will be the other team.
+
+    Args:
+        team_rosters: All teams and their players, as returned by
+            get_team_rosters().
+    """
+    if len(team_rosters.keys()) < 2:
+        raise exceptions.ServerError("Game has fewer than two teams in it")
+
+    team1, team2 = iter(team_rosters.keys())
+
+    if Team.RED in team_rosters:
+        players_matchmake_team1 = get_player_display_names(team_rosters[Team.RED])
+
+        # Pick whichever other team there is to be the other one.
+        other_team = team2 if team1 == Team.RED else team1
+        players_matchmake_team2 = get_player_display_names(team_rosters[other_team])
+    else:
+        # We shouldn't have an SM5 game where red doesn't play, but maybe
+        # somebody messed with the game editor.
+        players_matchmake_team1 = get_player_display_names(team_rosters[team1])
+        players_matchmake_team2 = get_player_display_names(team_rosters[team2])
+
+    return players_matchmake_team1, players_matchmake_team2

tests/helpers/environment.py

@@ -47,6 +47,11 @@ def get_green_team() -> Teams:
     return _GREEN_TEAM
 
 
+def get_blue_team() -> Teams:
+    assert _BLUE_TEAM
+    return _BLUE_TEAM
+
+
 async def setup_test_database():
     """Creates a test in-memory database using SQLite, connects Tortoise to it, and generates the schema.
 

tests/helpers/gamehelper_test.py

@@ -1,10 +1,14 @@
 import unittest
 
+from sanic import exceptions
+
 from db.sm5 import SM5Game, SM5Stats
-from helpers.gamehelper import get_players_from_team
+from db.types import Team
+from helpers.gamehelper import get_players_from_team, get_team_rosters, PlayerInfo, get_player_display_names, \
+    get_matchmaking_teams
 from helpers.statshelper import count_zaps, get_sm5_kd_ratio, get_sm5_score_components
 from tests.helpers.environment import setup_test_database, ENTITY_ID_1, ENTITY_ID_2, get_sm5_game_id, \
-    teardown_test_database, create_destroy_base_event, add_entity, get_red_team
+    teardown_test_database, create_destroy_base_event, add_entity, get_red_team, get_green_team, get_blue_team
 
 
 class TestGameHelper(unittest.IsolatedAsyncioTestCase):
@@ -32,6 +36,75 @@ async def test_get_players_from_team(self):
 
         self.assertCountEqual([player2, player3], players_in_team)
 
+    async def test_get_team_rosters(self):
+        entity1, entity_end1 = await add_entity(entity_id="@LoggedIn", name="Indy", team=get_red_team(), type="player")
+        entity2, entity_end2 = await add_entity(entity_id="Red Base", name="Red Base", team=get_red_team(), type="base")
+        entity3, entity_end3 = await add_entity(entity_id="@Member", name="Miles", team=get_green_team(), type="player")
+        entity4, entity_end4 = await add_entity(entity_id="NotLoggedIn", name="Bumblebee", team=get_red_team(), type="player")
+
+        roster = await get_team_rosters([entity1, entity2, entity3, entity4],
+                                        [entity_end1, entity_end2, entity_end3, entity_end4])
+
+        player1 = PlayerInfo(entity_start=entity1, entity_end=entity_end1, display_name="Indy")
+        # entity2 is not a player2 and should be ignored.
+        player3 = PlayerInfo(entity_start=entity3, entity_end=entity_end3, display_name="Miles")
+        player4 = PlayerInfo(entity_start=entity4, entity_end=entity_end4, display_name="NotLoggedIn")
+
+        self.assertDictEqual({
+            Team.RED: [player1, player4],
+            Team.GREEN: [player3]
+        }, roster)
+
+    async def test_get_player_display_names(self):
+        entity1, entity_end1 = await add_entity(entity_id="@LoggedIn", name="Indy", team=get_red_team(), type="player")
+        entity2, entity_end2 = await add_entity(entity_id="@Member", name="Miles", team=get_green_team(), type="player")
+
+        player1 = PlayerInfo(entity_start=entity1, entity_end=entity_end1, display_name="Indy")
+        player2 = PlayerInfo(entity_start=entity2, entity_end=entity_end2, display_name="Miles")
+
+        self.assertCountEqual(["Indy", "Miles"], get_player_display_names([player1, player2]))
+
+    async def test_get_matchmaking_teams(self):
+        entity1, entity_end1 = await add_entity(entity_id="@LoggedIn", name="Indy", team=get_red_team(), type="player")
+        entity2, entity_end2 = await add_entity(entity_id="@Member", name="Miles", team=get_green_team(), type="player")
+        entity3, entity_end3 = await add_entity(entity_id="NotLoggedIn", name="Bumblebee", team=get_red_team(), type="player")
+
+        roster = await get_team_rosters([entity1, entity2, entity3],
+                                        [entity_end1, entity_end2, entity_end3])
+
+        player_matchmaking_1, player_matchmaking_2 = get_matchmaking_teams(roster)
+
+        # Red team should be team 1.
+        self.assertCountEqual(["Indy", "NotLoggedIn"], player_matchmaking_1)
+        self.assertCountEqual(["Miles"], player_matchmaking_2)
+
+    async def test_get_matchmaking_teams_no_red_team(self):
+        entity1, entity_end1 = await add_entity(entity_id="@LoggedIn", name="Indy", team=get_blue_team(), type="player")
+        entity2, entity_end2 = await add_entity(entity_id="@Member", name="Miles", team=get_green_team(), type="player")
+        entity3, entity_end3 = await add_entity(entity_id="NotLoggedIn", name="Bumblebee", team=get_blue_team(), type="player")
+
+        roster = await get_team_rosters([entity1, entity2, entity3],
+                                        [entity_end1, entity_end2, entity_end3])
+
+        player_matchmaking_1, player_matchmaking_2 = get_matchmaking_teams(roster)
+
+        # The order of the teams is not defined.
+        if "Miles" in player_matchmaking_2:
+            self.assertCountEqual(["Indy", "NotLoggedIn"], player_matchmaking_1)
+            self.assertCountEqual(["Miles"], player_matchmaking_2)
+        else:
+            self.assertCountEqual(["Indy", "NotLoggedIn"], player_matchmaking_2)
+            self.assertCountEqual(["Miles"], player_matchmaking_1)
+
+    async def test_get_matchmaking_teams_one_team_only(self):
+        entity1, entity_end1 = await add_entity(entity_id="@LoggedIn", name="Indy", team=get_blue_team(), type="player")
+
+        roster = await get_team_rosters([entity1],
+                                        [entity_end1])
+
+        with self.assertRaises(exceptions.ServerError):
+            get_matchmaking_teams(roster)
+
 
 if __name__ == '__main__':
     unittest.main()