Procházet zdrojové kódy

Improving configurability of rollbot class

Kirk Trombley před 6 roky
rodič
revize
67e532ae5c
2 změnil soubory, kde provedl 14 přidání a 11 odebrání
  1. 4 2
      src/app.py
  2. 10 9
      src/rollbot.py

+ 4 - 2
src/app.py

@@ -32,12 +32,14 @@ dictConfig({
 app = Flask(__name__)
 app.config["PROPAGATE_EXCEPTIONS"] = True
 rollbot = Rollbot(
-    app.logger,
+    logger=app.logger,
     plugins=get_config("plugins"),
     aliases=get_config("aliases"),
     responses=get_config("responses"),
     lookup=BOTS_LOOKUP,
-    sleep_time=float(get_config("sleep_time"))
+    sleep_time=float(get_config("sleep_time")),
+    callback=post_message,
+    session_factory=db.session_scope
 )
 app.logger.info("Initializing database tables")
 db.init_db()

+ 10 - 9
src/rollbot.py

@@ -1,10 +1,9 @@
 import importlib
+import logging
 import time
 from threading import Thread
 
-import db
 from command_system import RollbotResponse, RollbotFailure, as_plugin
-from util import post_message
 
 
 def lift_response(call, response):
@@ -15,8 +14,10 @@ def lift_response(call, response):
 
 
 class Rollbot:
-    def __init__(self, logger, plugins={}, aliases={}, responses={}, lookup={}, sleep_time=0.0):
+    def __init__(self, logger=logging.getLogger(__name__), plugins={}, aliases={}, responses={}, lookup={}, sleep_time=0.0, session_factory=None, callback=None):
         self.logger = logger
+        self.session_factory = session_factory or (lambda: None)
+        self.callback = callback or (lambda txt, gid: self.logger.info(f"Responding to {gid} with {txt}"))
         self.commands = {}
         self.to_start = set()
         self.to_stop = set()
@@ -62,14 +63,14 @@ class Rollbot:
 
     def start_plugins(self):
         self.logger.info("Starting plugins")
-        with db.session_scope() as session:
+        with self.session_factory() as session:
             for cmd in self.to_start:
                 cmd.on_start(session)
         self.logger.info("Finished starting plugins")
 
     def shutdown_plugins(self):
         self.logger.info("Shutting down plugins")
-        with db.session_scope() as session:
+        with self.session_factory() as session:
             for cmd in self.to_stop:
                 cmd.on_shutdown(session)
         self.logger.info("Finished shutting down plugins")
@@ -85,7 +86,7 @@ class Rollbot:
             self.logger.warn(f"Message {message.message_id} had a command {message.command} that could not be run.")
             return RollbotResponse(message, failure=RollbotFailure.INVALID_COMMAND)
 
-        with db.session_scope() as session:
+        with self.session_factory() as session:
             response = plugin.on_command(session, message)
 
         if not response.is_success:
@@ -121,11 +122,11 @@ class Rollbot:
         bot_id = self.bot_lookup[msg.group_id]
         if response.is_success:
             if response.txt is not None:
-                post_message(response.txt, bot_id)
+                self.callback(response.txt, bot_id)
             if response.img is not None:
-                post_message(response.img, bot_id)
+                self.callback(response.img, bot_id)
         else:
-            post_message(response.failure_msg, bot_id)
+            self.callback(response.failure_msg, bot_id)
             self.logger.warning(f"Failed command response: {response}")
 
         t = time.time() - t