refactor and add tests

3 files changed, 379 insertions, 147 deletions

diff --git a/rater.py b/rater.py
new file mode 100644
index 0000000..2d07e03
--- /dev/null
+++ b/rater.py
@@ -0,0 +1,186 @@
+#!/usr/bin/python3
+import os
+import sys
+import logging
+
+import trueskill
+import mpmath
+mpmath.mp.dps = 50  # increase if FloatingPointError
+
+UNKNOWN_PLAYER_SIGMA = int(os.environ.get("UNKNOWN_PLAYER_SIGMA") or 500)
+TAU = float(os.environ.get("TAU") or 1000/100.0)
+
+# mapping from Tier (-1 - 30) to average skill tier points
+vst_points = {
+    -1: 1,
+    0: 1
+}
+for c in range(1, 12):
+    vst_points[c] = (109 + 1/11) * (c + 0.5)
+for c in range(1, 5):
+    vst_points[11 + c] = vst_points[11] + 50 * (c + 0.5)
+for c in range(1, 10):
+    vst_points[15 + c] = vst_points[15] + (66 + 2/3) * (c + 0.5)
+for c in range(1, 4):
+    vst_points[24 + c] = vst_points[24] + (133 + 1/3) * (c + 0.5)
+for c in range(1, 3):
+    vst_points[27 + c] = vst_points[27] + 200 * (c + 0.5)
+# ---
+
+env = trueskill.TrueSkill(
+    backend="mpmath",
+    mu=1500,
+    sigma=1000,
+    beta=10.0/30*3000,
+    tau=TAU,
+    draw_probability=0
+)
+
+"""
+Return a (mu, sigma) based on information known about a player.
+"""
+def get_trueskill_seed(player):
+    # fallback 1 - approximate by max(ranked, blitz) points
+    rank_points = None
+    if player.rank_points_ranked is not None and player.rank_points_ranked != 0:
+        rank_points = player.rank_points_ranked
+    if player.rank_points_blitz is not None and player.rank_points_blitz != 0:
+        if rank_points is None:
+            rank_points = player.rank_points_blitz
+        else:
+            if player.rank_points_blitz > rank_points:
+                rank_points = player.rank_points_blitz
+
+    if rank_points is not None:
+        sigma = UNKNOWN_PLAYER_SIGMA * (2.0/3.0)  # more accurate than skill tier = more trust
+        mu = float(rank_points) + sigma
+    else:
+        # fallback 2 - approximate rank points by skill tier
+        sigma = UNKNOWN_PLAYER_SIGMA
+        mu = vst_points[player.skill_tier] + sigma
+
+    return (mu, sigma)
+
+
+"""
+Mutate a match structure by updating TrueSkill values.
+Returns changed match object.
+"""
+def rate_match(match):
+    trueskill_column = None  # on player
+    if match.game_mode == "casual":  # mode names mapped by processor!!!
+        trueskill_column = "trueskill_casual"
+    if match.game_mode == "ranked":
+        trueskill_column = "trueskill_ranked"
+    if match.game_mode == "blitz":
+        trueskill_column = "trueskill_blitz"
+    if match.game_mode == "br":
+        trueskill_column = "trueskill_br"
+    if trueskill_column is None:
+        logger.info("got unsupported game mode %s", match.game_mode)
+        return
+
+    matchup_shared = []  # trueskill shared across all modes
+    matchup = []
+
+    anyAfk = False
+    if len(match.rosters) != 2 \
+            or len(match.rosters[0].participants) != 3 \
+            or len(match.rosters[1].participants) != 3:
+        logger.error("got an invalid matchup %s", match.api_id)
+        anyAfk = True
+
+    for participant in match.participants:
+        participant.participant_items[0].any_afk = False
+        if participant.went_afk == 1:
+            logger.info("got an afk matchup %s", match.api_id)
+            anyAfk = True
+            break
+
+    if anyAfk:
+        match.trueskill_quality = 0
+        for participant in match.participants:
+            participant.participant_items[0].any_afk = True
+        return
+
+    for roster in match.rosters:
+        team_shared = []
+        team = []
+        for participant in roster.participants:
+            player = participant.player[0]
+
+            # calculate TrueSkill shared across all modes = starting point for all modes
+            if player.trueskill_mu is not None:
+                mu_shared = player.trueskill_mu
+                sigma_shared = player.trueskill_sigma
+            else:
+                mu_shared, sigma_shared = get_trueskill_seed(player)
+
+            team_shared.append(env.create_rating(float(mu_shared), float(sigma_shared)))
+
+            # calculate queue specific TrueSkill
+            if getattr(player, trueskill_column + "_mu") is not None:
+                sigma = getattr(player, trueskill_column + "_sigma")
+                mu = getattr(player, trueskill_column + "_mu")
+            else:
+                # fallback 1 - approximate by shared trueskill
+                sigma = sigma_shared
+                mu = mu_shared
+
+            team.append(env.create_rating(float(mu), float(sigma)))
+
+
+        matchup_shared.append(team_shared)
+        matchup.append(team)
+
+    logger.info("got a valid matchup %s", match.api_id)
+
+    # store the fairness of the match using the shared TrueSkill
+    match.trueskill_quality = env.quality(matchup)
+
+    # shared TrueSkill (participant)
+    for team, roster in zip(env.rate(matchup_shared, ranks=[int(not r.winner) for r in match.rosters]),
+            match.rosters):
+        # lower rank is better = winner!
+        for rating, participant in zip(team, roster.participants):
+            player = participant.player[0]
+            # delta = current - pre
+            if player.trueskill_mu is not None:
+                participant.trueskill_delta = (float(rating.mu) - float(rating.sigma)) - (float(player.trueskill_mu) - float(player.trueskill_sigma))
+            else:
+                participant.trueskill_delta = 0
+            player.trueskill_mu = rating.mu
+            participant.trueskill_mu = rating.mu
+            player.trueskill_sigma = rating.sigma
+            participant.trueskill_sigma = rating.sigma
+
+    # queue specific TrueSkill (participant_items)
+    # no delta
+    for team, roster in zip(env.rate(matchup, ranks=[int(not r.winner) for r in match.rosters]),
+            match.rosters):
+        for rating, participant in zip(team, roster.participants):
+            player = participant.player[0]
+            pi = participant.participant_items[0]
+            setattr(player, trueskill_column + "_mu", rating.mu)
+            setattr(pi, trueskill_column + "_mu", rating.mu)
+            setattr(player, trueskill_column + "_sigma", rating.sigma)
+            setattr(pi, trueskill_column + "_sigma", rating.sigma)
+
+
+# TODO share this between the two classes
+# log errs to stderr, debug to stdout
+class InfoFilter(logging.Filter):
+    def filter(self, record):
+        return record.levelno in (logging.DEBUG, logging.INFO)
+
+logger = logging.getLogger("__name__")
+logger.setLevel(logging.INFO)
+
+h1 = logging.StreamHandler(sys.stdout)
+h1.setLevel(logging.INFO)
+h1.addFilter(InfoFilter())
+logger.addHandler(h1)
+
+h2 = logging.StreamHandler(sys.stderr)
+h2.setLevel(logging.WARNING)
+logger.addHandler(h2)
diff --git a/worker.py b/worker.py
index 520d9fa..4771646 100644
--- a/worker.py
+++ b/worker.py
@@ -9,9 +9,8 @@ from sqlalchemy.ext.automap import automap_base
 from sqlalchemy import create_engine
 
 import pika
-import trueskill
-import mpmath
-mpmath.mp.dps = 50  # increase if FloatingPointError
+
+import rater
 
 
 RABBITMQ_URI = os.environ.get("RABBITMQ_URI") or "amqp://localhost"
@@ -26,26 +25,6 @@ DOTELESUCKMATCH = os.environ.get("DOTELESUCKMATCH") == "true"
 TELESUCK_QUEUE = os.environ.get("TELESUCK_QUEUE") or "telesuck"
 DOSEWMATCH = os.environ.get("DOSEWMATCH") == "true"
 SEW_QUEUE = os.environ.get("SEW_QUEUE") or "sew"
-UNKNOWN_PLAYER_SIGMA = int(os.environ.get("UNKNOWN_PLAYER_SIGMA") or 500)
-TAU = float(os.environ.get("TAU") or 1000/100.0)
-
-# mapping from Tier (-1 - 30) to average skill tier points
-vst_points = {
-    -1: 1,
-    0: 1
-}
-for c in range(1, 12):
-    vst_points[c] = (109 + 1/11) * (c + 0.5)
-for c in range(1, 5):
-    vst_points[11 + c] = vst_points[11] + 50 * (c + 0.5)
-for c in range(1, 10):
-    vst_points[15 + c] = vst_points[15] + (66 + 2/3) * (c + 0.5)
-for c in range(1, 4):
-    vst_points[24 + c] = vst_points[24] + (133 + 1/3) * (c + 0.5)
-for c in range(1, 3):
-    vst_points[27 + c] = vst_points[27] + 200 * (c + 0.5)
-# ---
-
 
 # ORM definitions
 Match = Asset = Roster = Participant = ParticipantStats = ParticipantItems = Player = None
@@ -194,14 +173,6 @@ def process():
 
     db = Session()
     try:
-        env = trueskill.TrueSkill(
-            backend="mpmath",
-            mu=1500,
-            sigma=1000,
-            beta=10.0/30*3000,
-            tau=TAU,
-            draw_probability=0
-        )
         for match in db.query(Match).order_by(Match.created_at.asc()).options(
             load_only("api_id", "game_mode")
             .selectinload(Match.rosters)
@@ -211,128 +182,14 @@ def process():
                            "player_api_id", "skill_tier", "went_afk")
                 .selectinload(Participant.player)
                     .load_only("api_id",
-                               "rank_points_ranked",
+                               "rank_points_ranked", "rank_points_blitz",
                                "trueskill_sigma", "trueskill_mu",
                                "trueskill_casual_sigma", "trueskill_casual_mu",
                                "trueskill_ranked_sigma", "trueskill_ranked_mu",
                                "trueskill_blitz_sigma", "trueskill_blitz_mu",
                                "trueskill_br_sigma", "trueskill_br_mu")
          ).filter(Match.api_id.in_(ids)).yield_per(CHUNKSIZE):
-            trueskill_column = None  # on player
-            if match.game_mode == "casual":  # mode names mapped by processor!!!
-                trueskill_column = "trueskill_casual"
-            if match.game_mode == "ranked":
-                trueskill_column = "trueskill_ranked"
-            if match.game_mode == "blitz":
-                trueskill_column = "trueskill_blitz"
-            if match.game_mode == "br":
-                trueskill_column = "trueskill_br"
-            if trueskill_column is None:
-                logger.info("got unsupported game mode %s", match.game_mode)
-                continue
-
-            matchup_shared = []  # trueskill shared across all modes
-            matchup = []
-
-            anyAfk = False
-            if len(match.rosters) != 2 \
-                    or len(match.rosters[0].participants) != 3 \
-                    or len(match.rosters[1].participants) != 3:
-                logger.error("got an invalid matchup %s", match.api_id)
-                anyAfk = True
-
-            for participant in match.participants:
-                participant.participant_items[0].any_afk = False
-                if participant.went_afk == 1:
-                    logger.info("got an afk matchup %s", match.api_id)
-                    anyAfk = True
-                    break
-
-            if anyAfk:
-                match.trueskill_quality = 0
-                for participant in match.participants:
-                    participant.participant_items[0].any_afk = True
-                continue
-
-            for roster in match.rosters:
-                team_shared = []
-                team = []
-                for participant in roster.participants:
-                    player = participant.player[0]
-
-                    # calculate TrueSkill shared across all modes = starting point for all modes
-                    if player.trueskill_mu is not None:
-                        mu_shared = player.trueskill_mu
-                        sigma_shared = player.trueskill_sigma
-                    else:
-                        # fallback 1 - approximate by max(ranked, blitz) points
-                        rank_points = None
-                        if player.rank_points_ranked is not None and player.rank_points_ranked != 0:
-                            rank_points = player.rank_points_ranked
-                        if player.rank_points_blitz is not None and player.rank_points_blitz != 0:
-                            if rank_points is None:
-                                rank_points = player.rank_points_blitz
-                            else:
-                                if player.rank_points_blitz > rank_points:
-                                    rank_points = player.rank_points_blitz
-
-                        if rank_points is not None:
-                            sigma_shared = UNKNOWN_PLAYER_SIGMA * (2.0/3.0)  # more accurate than skill tier = more trust
-                            mu_shared = float(rank_points) + sigma_shared
-                        else:
-                            # fallback 2 - approximate rank points by skill tier
-                            sigma_shared = UNKNOWN_PLAYER_SIGMA
-                            mu_shared = vst_points[participant.skill_tier] + sigma_shared
-
-                    team_shared.append(env.create_rating(float(mu_shared), float(sigma_shared)))
-
-                    # calculate queue specific TrueSkill
-                    if getattr(player, trueskill_column + "_mu") is not None:
-                        sigma = getattr(player, trueskill_column + "_sigma")
-                        mu = getattr(player, trueskill_column + "_mu")
-                    else:
-                        # fallback 1 - approximate by shared trueskill
-                        sigma = sigma_shared
-                        mu = mu_shared
-
-                    team.append(env.create_rating(float(mu), float(sigma)))
-
-
-                matchup_shared.append(team_shared)
-                matchup.append(team)
-
-            logger.info("got a valid matchup %s", match.api_id)
-
-            # store the fairness of the match using the shared TrueSkill
-            match.trueskill_quality = env.quality(matchup)
-
-            # shared TrueSkill (participant)
-            for team, roster in zip(env.rate(matchup_shared, ranks=[int(not r.winner) for r in match.rosters]),
-                                    match.rosters):
-                # lower rank is better = winner!
-                for rating, participant in zip(team, roster.participants):
-                    player = participant.player[0]
-                    # delta = current - pre
-                    if player.trueskill_mu is not None:
-                        participant.trueskill_delta = (float(rating.mu) - float(rating.sigma)) - (float(player.trueskill_mu) - float(player.trueskill_sigma))
-                    else:
-                        participant.trueskill_delta = 0
-                    player.trueskill_mu = rating.mu
-                    participant.trueskill_mu = rating.mu
-                    player.trueskill_sigma = rating.sigma
-                    participant.trueskill_sigma = rating.sigma
-
-            # queue specific TrueSkill (participant_items)
-            # no delta
-            for team, roster in zip(env.rate(matchup, ranks=[int(not r.winner) for r in match.rosters]),
-                                    match.rosters):
-                for rating, participant in zip(team, roster.participants):
-                    player = participant.player[0]
-                    pi = participant.participant_items[0]
-                    setattr(player, trueskill_column + "_mu", rating.mu)
-                    setattr(pi, trueskill_column + "_mu", rating.mu)
-                    setattr(player, trueskill_column + "_sigma", rating.sigma)
-                    setattr(pi, trueskill_column + "_sigma", rating.sigma)
+            match = rater.rate_match(match)
 
         db.commit()
     except:
diff --git a/worker_test.py b/worker_test.py
new file mode 100644
index 0000000..22bedcb
--- /dev/null
+++ b/worker_test.py
@@ -0,0 +1,189 @@
+#!/usr/bin/python
+
+import pytest
+import rater
+
+class Match:
+    def __init__(self, game_mode, rosters):
+        self.api_id = ""
+        self.game_mode = game_mode
+        self.rosters = rosters
+        self.participants = [p for r in rosters for p in r.participants]
+
+class Roster:
+    def __init__(self, winner, participants):
+        self.api_id = ""
+        self.winner = winner
+        self.participants = participants
+
+class Participant:
+    def __init__(self, skill_tier, went_afk,
+                 trueskill_mu, trueskill_sigma,
+                 participant_items, player):
+        self.api_id = ""
+        self.skill_tier = skill_tier
+        self.went_afk = went_afk
+        self.participant_items = [participant_items]
+        self.player = [player]
+
+class ParticipantItems:
+    def __init__(self, trueskill_casual_mu, trueskill_casual_sigma,
+                 trueskill_ranked_mu, trueskill_ranked_sigma,
+                 trueskill_blitz_mu, trueskill_blitz_sigma,
+                 trueskill_br_mu, trueskill_br_sigma):
+        self.api_id = ""
+        self.trueskill_casual_mu = trueskill_casual_mu
+        self.trueskill_casual_sigma = trueskill_casual_sigma
+        self.trueskill_ranked_mu = trueskill_ranked_mu
+        self.trueskill_ranked_sigma = trueskill_ranked_sigma
+        self.trueskill_blitz_mu = trueskill_blitz_mu
+        self.trueskill_blitz_sigma = trueskill_blitz_sigma
+        self.any_afk = False
+
+class Player:
+    def __init__(self, skill_tier,
+                 rank_points_ranked,
+                 rank_points_blitz,
+                 trueskill_mu, trueskill_sigma,
+                 trueskill_casual_mu, trueskill_casual_sigma,
+                 trueskill_ranked_mu, trueskill_ranked_sigma,
+                 trueskill_blitz_mu, trueskill_blitz_sigma,
+                 trueskill_br_mu, trueskill_br_sigma):
+        self.api_id = ""
+        self.skill_tier = skill_tier
+        self.rank_points_ranked = rank_points_ranked
+        self.rank_points_blitz = rank_points_blitz
+        self.trueskill_mu = trueskill_mu
+        self.trueskill_sigma = trueskill_sigma
+        self.trueskill_casual_mu = trueskill_casual_mu
+        self.trueskill_casual_sigma = trueskill_casual_sigma
+        self.trueskill_ranked_mu = trueskill_ranked_mu
+        self.trueskill_ranked_sigma = trueskill_ranked_sigma
+        self.trueskill_blitz_mu = trueskill_blitz_mu
+        self.trueskill_blitz_sigma = trueskill_blitz_sigma
+
+
+class TestRater:
+    def test_get_trueskill_seed(self):
+        # test approx. by skill tier
+        player = Player(15, None, None,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None)
+        mu, sigma = rater.get_trueskill_seed(player)
+        assert 1300 < mu - sigma < 1700
+
+        # test approx. by rank points
+        player = Player(0, 2500, None,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None)
+        mu, sigma = rater.get_trueskill_seed(player)
+        assert mu - sigma == 2500
+
+        player = Player(0, 2500, 100,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None)
+        mu, sigma = rater.get_trueskill_seed(player)
+        assert mu - sigma == 2500
+
+        player = Player(0, 100, 2500,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None)
+        mu, sigma = rater.get_trueskill_seed(player)
+        assert mu - sigma == 2500
+
+        player = Player(0, None, 2500,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None,
+                        None, None)
+        mu, sigma = rater.get_trueskill_seed(player)
+        assert mu - sigma == 2500
+
+    def test_rate_match(self):
+        # new user, only skill tier
+        g_player = lambda: Player(15, None, None,
+                                None, None,
+                                None, None,
+                                None, None,
+                                None, None,
+                                None, None)
+        g_participant_items = lambda: ParticipantItems(None, None,
+                                                     None, None,
+                                                     None, None,
+                                                     None, None)
+        g_participant = lambda: Participant(0, False,
+                                          None, None,
+                                          g_participant_items(), g_player())
+        roster = Roster(True, [g_participant()]*3)
+        roster2 = Roster(False, [g_participant()]*3)
+        match = Match("ranked", [roster, roster2])
+
+        rater.rate_match(match)
+
+        assert match.rosters[0].participants[0].player[0].trueskill_mu is not None
+        assert match.rosters[0].participants[0].player[0].trueskill_ranked_mu is not None
+        assert match.rosters[0].participants[0].player[0].trueskill_ranked_sigma < match.rosters[0].participants[0].player[0].trueskill_ranked_mu
+        assert 500 < match.rosters[0].participants[0].player[0].trueskill_ranked_mu < 2500
+        assert match.rosters[0].participants[0].player[0].trueskill_casual_mu is None
+        assert match.rosters[0].participants[0].player[0].trueskill_mu > match.rosters[1].participants[0].player[0].trueskill_mu
+        assert match.rosters[0].participants[0].player[0].trueskill_ranked_mu > match.rosters[1].participants[0].player[0].trueskill_ranked_mu
+
+    def test_rate_match_returning(self):
+        # returning user
+        g_player = lambda: Player(None, None, None,
+                                  2000, 100,
+                                  None, None,
+                                  None, None,
+                                  None, None,
+                                  None, None)
+        g_participant_items = lambda: ParticipantItems(None, None,
+                                                       None, None,
+                                                       None, None,
+                                                       None, None)
+        g_participant = lambda: Participant(0, False,
+                                            None, None,
+                                            g_participant_items(), g_player())
+        roster = Roster(True, [g_participant()]*3)
+        roster2 = Roster(False, [g_participant()]*3)
+        match = Match("ranked", [roster, roster2])
+
+        rater.rate_match(match)
+
+        assert 1800 < match.rosters[0].participants[0].player[0].trueskill_ranked_mu < 2200
+
+    def test_rate_match_afk(self):
+        # afk
+        g_player = lambda: Player(None, None, None,
+                                  None, None,
+                                  None, None,
+                                  None, None,
+                                  None, None,
+                                  None, None)
+        g_participant_items = lambda: ParticipantItems(None, None,
+                                                       None, None,
+                                                       None, None,
+                                                       None, None)
+        g_participant = lambda: Participant(0, True,
+                                            None, None,
+                                            g_participant_items(), g_player())
+        roster = Roster(True, [g_participant()]*3)
+        roster2 = Roster(False, [g_participant()]*3)
+        match = Match("ranked", [roster, roster2])
+
+        rater.rate_match(match)
+
+        assert match.rosters[0].participants[0].player[0].trueskill_mu is None
+        assert match.rosters[0].participants[0].participant_items[0].any_afk == True