From 285a946a1f1ce7874875e882a0e44836ede9ac71 Mon Sep 17 00:00:00 2001
From: PIG208 <359101898@qq.com>
Date: Thu, 22 Jul 2021 13:04:54 +0800
Subject: [PATCH] bot_server: Add support for running botserver from bots
 registry.

---
 zulip_botserver/tests/test_server.py      | 32 +++++++++++++++++++++++
 zulip_botserver/zulip_botserver/server.py | 20 +++++++-------
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/zulip_botserver/tests/test_server.py b/zulip_botserver/tests/test_server.py
index 1b1b6fe..59367f7 100644
--- a/zulip_botserver/tests/test_server.py
+++ b/zulip_botserver/tests/test_server.py
@@ -8,6 +8,7 @@ from types import ModuleType
 from typing import Any, Dict
 from unittest import mock
 
+from zulip_bots.finder import metadata
 from zulip_bots.lib import BotHandler
 from zulip_botserver import server
 from zulip_botserver.input_parameters import parse_args
@@ -273,6 +274,37 @@ class BotServerTests(BotServerTestCase):
             ).as_posix()
             module = server.load_lib_modules([path])[path]
 
+    @mock.patch("zulip_botserver.server.app")
+    @mock.patch("sys.argv", ["zulip-botserver", "--config-file", "/foo/bar/baz.conf"])
+    def test_load_from_registry(self, mock_app: mock.Mock) -> None:
+        packaged_bot_module = mock.MagicMock(__version__="1.0.0", __file__="asd")
+        packaged_bot_entrypoint = metadata.EntryPoint(
+            "packaged_bot", "module_name", "zulip_bots.registry"
+        )
+        bots_config = {
+            "packaged_bot": {
+                "email": "packaged-bot@zulip.com",
+                "key": "value",
+                "site": "http://localhost",
+                "token": "abcd1234",
+            }
+        }
+
+        with mock.patch(
+            "zulip_botserver.server.read_config_file", return_value=bots_config
+        ), mock.patch("zulip_botserver.server.lib.ExternalBotHandler", new=mock.Mock()), mock.patch(
+            "zulip_bots.finder.metadata.EntryPoint.load",
+            return_value=packaged_bot_module,
+        ), mock.patch(
+            "zulip_bots.finder.metadata.entry_points",
+            return_value=(packaged_bot_entrypoint,),
+        ):
+            server.main()
+
+        mock_app.config.__setitem__.assert_any_call(
+            "BOTS_LIB_MODULES", {"packaged_bot": packaged_bot_module}
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/zulip_botserver/zulip_botserver/server.py b/zulip_botserver/zulip_botserver/server.py
index b97f63a..17d0511 100644
--- a/zulip_botserver/zulip_botserver/server.py
+++ b/zulip_botserver/zulip_botserver/server.py
@@ -16,7 +16,7 @@ from werkzeug.exceptions import BadRequest, Unauthorized
 
 from zulip import Client
 from zulip_bots import lib
-from zulip_bots.finder import import_module_from_source
+from zulip_bots.finder import import_module_from_source, import_module_from_zulip_bot_registry
 from zulip_botserver.input_parameters import parse_args
 
 
@@ -122,15 +122,17 @@ def load_lib_modules(available_bots: List[str]) -> Dict[str, ModuleType]:
                 lib_module = import_module(module_name)
             bots_lib_module[bot] = lib_module
         except ImportError:
-            error_message = (
-                'Error: Bot "{}" doesn\'t exist. Please make sure '
-                "you have set up the botserverrc file correctly.\n".format(bot)
-            )
-            if bot == "api":
-                error_message += (
-                    "Did you forget to specify the bot you want to run with -b <botname> ?"
+            _, bots_lib_module[bot] = import_module_from_zulip_bot_registry(bot)
+            if bots_lib_module[bot] is None:
+                error_message = (
+                    f'Error: Bot "{bot}" doesn\'t exist. Please make sure '
+                    "you have set up the botserverrc file correctly.\n"
                 )
-            sys.exit(error_message)
+                if bot == "api":
+                    error_message += (
+                        "Did you forget to specify the bot you want to run with -b <botname> ?"
+                    )
+                sys.exit(error_message)
     return bots_lib_module