allow subclasses to pass protocol and bus objects (#76)

JarbasHiveMind · Sep 28, 2023 · a5855dd · a5855dd
1 parent abf6b8b
commit a5855dd
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/hivemind_core/service.py b/hivemind_core/service.py
@@ -166,7 +166,9 @@ def __init__(self,
                  ready_hook: Callable = on_ready,
                  error_hook: Callable = on_error,
                  stopping_hook: Callable = on_stopping,
-                 websocket_config: Optional[Dict[str, Any]] = None):
+                 websocket_config: Optional[Dict[str, Any]] = None,
+                 protocol=HiveMindListenerProtocol,
+                 bus=None):
 
         websocket_config = websocket_config or \
                 Configuration().get('hivemind_websocket', {})
@@ -175,10 +177,13 @@ def __init__(self,
                                       on_ready=ready_hook,
                                       on_error=error_hook,
                                       on_stopping=stopping_hook)
-
-        self.bus = MessageBusClient(emitter=EventEmitter())
-        self.bus.run_in_thread()
-        self.bus.connected_event.wait()
+        self._proto = protocol
+        if bus:
+            self.bus = bus
+        else:
+            self.bus = MessageBusClient(emitter=EventEmitter())
+            self.bus.run_in_thread()
+            self.bus.connected_event.wait()
 
         self.status = ProcessStatus('HiveMind', callback_map=callbacks)
         self.host = websocket_config.get('host') or "0.0.0.0"
@@ -205,7 +210,7 @@ def run(self):
         asyncio.set_event_loop_policy(AnyThreadEventLoopPolicy())
         loop = ioloop.IOLoop.current()
 
-        self.protocol = HiveMindListenerProtocol(loop=loop)
+        self.protocol = self._proto(loop=loop)
         self.protocol.bind(MessageBusEventHandler, self.bus)
         self.status.bind(self.bus)
         self.status.set_started()