Merge pull request #496 from linuxdaemon/gonzobot+refactor-duck-score

Refactor duckhunt .killers/.friends commands

plugins/duckhunt.py

@@ -4,13 +4,14 @@
 from threading import Lock
 from time import time, sleep
 
-from sqlalchemy import Table, Column, String, Integer, PrimaryKeyConstraint, desc, Boolean, and_
+from sqlalchemy import Boolean, Column, Integer, PrimaryKeyConstraint, String, Table, and_, desc
 from sqlalchemy.sql import select
 
 from cloudbot import hook
 from cloudbot.event import EventType
 from cloudbot.util import database
 from cloudbot.util.formatting import pluralize_auto, truncate
+from cloudbot.util.func_utils import call_with_args
 
 duck_tail = "・゜゜・。。・゜゜"
 duck = ["\\_o< ", "\\_O< ", "\\_0< ", "\\_\u00f6< ", "\\_\u00f8< ", "\\_\u00f3< "]
@@ -584,105 +585,139 @@ def top_list(prefix, data, join_char=' • '):
     )
 
 
+def get_scores(db, score_type, network, chan=None):
+    clause = table.c.network == network
+    if chan is not None:
+        clause = and_(clause, table.c.chan == chan.lower())
+
+    query = select([table.c.name, table.c[score_type]], clause) \
+        .order_by(desc(table.c[score_type]))
+
+    scores = db.execute(query).fetchall()
+    return scores
+
+
+class ScoreType:
+    def __init__(self, name, column_name, noun, verb):
+        self.name = name
+        self.column_name = column_name
+        self.noun = noun
+        self.verb = verb
+
+
+def get_channel_scores(db, score_type: ScoreType, conn, chan):
+    scores_dict = defaultdict(int)
+    scores = get_scores(db, score_type.column_name, conn.name, chan)
+    if not scores:
+        return None
+
+    for row in scores:
+        if row[1] == 0:
+            continue
+
+        scores_dict[row[0]] += row[1]
+
+    return scores_dict
+
+
+def _get_global_scores(db, score_type: ScoreType, conn):
+    scores_dict = defaultdict(int)
+    chancount = defaultdict(int)
+    scores = get_scores(db, score_type.column_name, conn.name)
+    if not scores:
+        return None, None
+
+    for row in scores:
+        if row[1] == 0:
+            continue
+
+        chancount[row[0]] += 1
+        scores_dict[row[0]] += row[1]
+
+    return scores_dict, chancount
+
+
+def get_global_scores(db, score_type: ScoreType, conn):
+    return _get_global_scores(db, score_type, conn)[0]
+
+
+def get_average_scores(db, score_type: ScoreType, conn):
+    scores_dict, chancount = _get_global_scores(db, score_type, conn)
+    if not scores_dict:
+        return None
+
+    for k, v in scores_dict.items():
+        scores_dict[k] = int(v / chancount[k])
+
+    return scores_dict
+
+
+SCORE_TYPES = {
+    'friend': ScoreType('befriend', 'befriend', 'friend', 'friended'),
+    'killer': ScoreType('killer', 'shot', 'killer', 'killed'),
+}
+
+DISPLAY_FUNCS = {
+    'average': get_average_scores,
+    'global': get_global_scores,
+    None: get_channel_scores,
+}
+
+
+def display_scores(score_type: ScoreType, text, chan, conn, db):
+    if is_opt_out(conn.name, chan):
+        return
+
+    global_pfx = "Duck {noun} scores across the network: ".format(
+        noun=score_type.noun
+    )
+    chan_pfx = "Duck {noun} scores in {chan}: ".format(
+        noun=score_type.noun, chan=chan
+    )
+    no_ducks = "It appears no one has {verb} any ducks yet."
+
+    out = global_pfx if text else chan_pfx
+
+    func = DISPLAY_FUNCS[text.lower() or None]
+    scores_dict = call_with_args(func, {
+        'db': db,
+        'score_type': score_type,
+        'conn': conn,
+        'chan': chan,
+    })
+
+    if not scores_dict:
+        return no_ducks
+
+    return top_list(out, scores_dict.items())
+
+
 @hook.command("friends", autohelp=False)
 def friends(text, chan, conn, db):
-    """[{global|average}] - Prints a list of the top duck friends in the channel, if 'global' is specified all channels
-    in the database are included.
+    """[{global|average}] - Prints a list of the top duck friends in the
+    channel, if 'global' is specified all channels in the database are
+    included.
 
     :type text: str
     :type chan: str
     :type conn: cloudbot.client.Client
     :type db: sqlalchemy.orm.Session
     """
-    if is_opt_out(conn.name, chan):
-        return
-
-    friends_dict = defaultdict(int)
-    chancount = defaultdict(int)
-    if text.lower() in ('global', 'average'):
-        out = "Duck friend scores across the network: "
-        scores = db.execute(select([table.c.name, table.c.befriend]) \
-                            .where(table.c.network == conn.name) \
-                            .order_by(desc(table.c.befriend)))
-        if scores:
-            for row in scores:
-                if row[1] == 0:
-                    continue
-
-                chancount[row[0]] += 1
-                friends_dict[row[0]] += row[1]
-
-            if text.lower() == 'average':
-                for k, v in friends_dict.items():
-                    friends_dict[k] = int(v / chancount[k])
-        else:
-            return "it appears no on has friended any ducks yet."
-    else:
-        out = "Duck friend scores in {}: ".format(chan)
-        scores = db.execute(select([table.c.name, table.c.befriend]) \
-                            .where(table.c.network == conn.name) \
-                            .where(table.c.chan == chan.lower()) \
-                            .order_by(desc(table.c.befriend)))
-        if scores:
-            for row in scores:
-                if row[1] == 0:
-                    continue
-                friends_dict[row[0]] += row[1]
-        else:
-            return "it appears no on has friended any ducks yet."
-
-    return top_list(out, friends_dict.items())
+    return display_scores(SCORE_TYPES['friend'], text, chan, conn, db)
 
 
 @hook.command("killers", autohelp=False)
 def killers(text, chan, conn, db):
-    """[{global|average}] - Prints a list of the top duck killers in the channel, if 'global' is specified all channels
-    in the database are included.
+    """[{global|average}] - Prints a list of the top duck killers in the
+    channel, if 'global' is specified all channels in the database are
+    included.
 
     :type text: str
     :type chan: str
     :type conn: cloudbot.client.Client
     :type db: sqlalchemy.orm.Session
     """
-    if is_opt_out(conn.name, chan):
-        return
-
-    killers_dict = defaultdict(int)
-    chancount = defaultdict(int)
-    if text.lower() in ('global', 'average'):
-        out = "Duck killer scores across the network: "
-        scores = db.execute(select([table.c.name, table.c.shot]) \
-                            .where(table.c.network == conn.name) \
-                            .order_by(desc(table.c.shot)))
-        if scores:
-            for row in scores:
-                if row[1] == 0:
-                    continue
-
-                chancount[row[0]] += 1
-                killers_dict[row[0]] += row[1]
-
-            if text.lower() == 'average':
-                for k, v in killers_dict.items():
-                    killers_dict[k] = int(v / chancount[k])
-        else:
-            return "it appears no on has killed any ducks yet."
-    else:
-        out = "Duck killer scores in {}: ".format(chan)
-        scores = db.execute(select([table.c.name, table.c.shot]) \
-                            .where(table.c.network == conn.name) \
-                            .where(table.c.chan == chan.lower()) \
-                            .order_by(desc(table.c.shot)))
-        if scores:
-            for row in scores:
-                if row[1] == 0:
-                    continue
-
-                killers_dict[row[0]] += row[1]
-        else:
-            return "it appears no on has killed any ducks yet."
-
-    return top_list(out, killers_dict.items())
+    return display_scores(SCORE_TYPES['killer'], text, chan, conn, db)
 
 
 @hook.command("duckforgive", permissions=["op", "ignore"])

tests/plugin_tests/test_duckhunt.py

@@ -0,0 +1,116 @@
+import importlib
+
+import pytest
+from sqlalchemy import create_engine
+from sqlalchemy.orm import scoped_session, sessionmaker
+
+from plugins.duckhunt import top_list
+
+
+class MockDB:
+    def __init__(self):
+        self.engine = create_engine('sqlite:///:memory:')
+        self.session = scoped_session(sessionmaker(self.engine))
+
+
+@pytest.fixture()
+def mock_db():
+    return MockDB()
+
+
+@pytest.mark.parametrize('prefix,items,result', [
+    [
+        'Duck friend scores in #TestChannel: ',
+        {
+            'testuser': 5,
+            'testuser1': 1,
+        },
+        'Duck friend scores in #TestChannel: '
+        '\x02t\u200bestuser\x02: 5 • \x02t\u200bestuser1\x02: 1'
+    ],
+])
+def test_top_list(prefix, items, result):
+    assert top_list(prefix, items.items()) == result
+
+
+def test_display_scores(mock_db):
+    from cloudbot.util.database import metadata
+    metadata.bind = mock_db.engine
+    from plugins import duckhunt
+
+    importlib.reload(duckhunt)
+
+    metadata.create_all(checkfirst=True)
+
+    session = mock_db.session()
+
+    class Conn:
+        name = 'TestConn'
+
+    conn = Conn()
+
+    chan = '#TestChannel'
+    chan1 = '#TestChannel1'
+
+    duckhunt.update_score('TestUser', chan, session, conn, 5, 4)
+    duckhunt.update_score('TestUser1', chan, session, conn, 1, 7)
+    duckhunt.update_score('OtherUser', chan1, session, conn, 9, 2)
+
+    expected_testchan_friend_scores = {'testuser': 4, 'testuser1': 7}
+
+    actual_testchan_friend_scores = duckhunt.get_channel_scores(
+        session, duckhunt.SCORE_TYPES['friend'],
+        conn, chan
+    )
+
+    assert actual_testchan_friend_scores == expected_testchan_friend_scores
+
+    chan_friends = (
+        'Duck friend scores in #TestChannel: '
+        '\x02t\u200bestuser1\x02: 7 • \x02t\u200bestuser\x02: 4'
+    )
+
+    chan_kills = (
+        'Duck killer scores in #TestChannel: '
+        '\x02t\u200bestuser\x02: 5 • \x02t\u200bestuser1\x02: 1'
+    )
+
+    global_friends = (
+        'Duck friend scores across the network: '
+        '\x02t\u200bestuser1\x02: 7'
+        ' • \x02t\u200bestuser\x02: 4'
+        ' • \x02o\u200btheruser\x02: 2'
+    )
+
+    global_kills = (
+        'Duck killer scores across the network: '
+        '\x02o\u200btheruser\x02: 9'
+        ' • \x02t\u200bestuser\x02: 5'
+        ' • \x02t\u200bestuser1\x02: 1'
+    )
+
+    average_friends = (
+        'Duck friend scores across the network: '
+        '\x02t\u200bestuser1\x02: 7'
+        ' • \x02t\u200bestuser\x02: 4'
+        ' • \x02o\u200btheruser\x02: 2'
+    )
+
+    average_kills = (
+        'Duck killer scores across the network: '
+        '\x02o\u200btheruser\x02: 9'
+        ' • \x02t\u200bestuser\x02: 5'
+        ' • \x02t\u200bestuser1\x02: 1'
+    )
+
+    assert duckhunt.friends('', chan, conn, session) == chan_friends
+
+    assert duckhunt.killers('', chan, conn, session) == chan_kills
+
+    assert duckhunt.friends('global', chan, conn, session) == global_friends
+
+    assert duckhunt.killers('global', chan, conn, session) == global_kills
+
+    assert duckhunt.friends('average', chan, conn, session) == average_friends
+
+    assert duckhunt.killers('average', chan, conn, session) == average_kills