From e5239c5c5439146f5ebcc61239188280d6887e25 Mon Sep 17 00:00:00 2001
From: novokrest <novokrest013@gmail.com>
Date: Sat, 5 May 2018 09:10:36 +0300
Subject: [PATCH] lib: Add ExternalBotHandler methods for file uploading.

Add methods ExternalBotHandler.upload_file() and
ExternalBotHandler.upload_file_from_path() for uploading
local files to Zulip server.

Fixes #351
---
 zulip_bots/zulip_bots/lib.py            |  9 +++++++
 zulip_bots/zulip_bots/tests/test_lib.py | 34 +++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/zulip_bots/zulip_bots/lib.py b/zulip_bots/zulip_bots/lib.py
index 875300e..4221248 100644
--- a/zulip_bots/zulip_bots/lib.py
+++ b/zulip_bots/zulip_bots/lib.py
@@ -211,6 +211,15 @@ class ExternalBotHandler(object):
 
         return dict(config_parser.items(bot_name))
 
+    def upload_file_from_path(self, file_path: str) -> Dict[str, Any]:
+        with open(file_path, 'rb') as file:
+            return self.upload_file(file)
+
+    def upload_file(self, file: IO[Any]) -> Dict[str, Any]:
+        if not self._rate_limit.is_legal():
+            self._rate_limit.show_error_and_exit()
+        return self._client.upload_file(file)
+
     def open(self, filepath: str) -> IO[str]:
         filepath = os.path.normpath(filepath)
         abs_filepath = os.path.join(self._root_dir, filepath)
diff --git a/zulip_bots/zulip_bots/tests/test_lib.py b/zulip_bots/zulip_bots/tests/test_lib.py
index f3c9c0f..9af194a 100644
--- a/zulip_bots/zulip_bots/tests/test_lib.py
+++ b/zulip_bots/zulip_bots/tests/test_lib.py
@@ -6,6 +6,8 @@ from zulip_bots.lib import (
     run_message_handler_for_bot,
 )
 
+import io
+
 class FakeClient:
     def __init__(self, *args, **kwargs):
         self.storage = dict()
@@ -36,6 +38,9 @@ class FakeClient:
             result='success',
         )
 
+    def upload_file(self, file):
+        pass
+
 class FakeBotHandler:
     def usage(self):
         return '''
@@ -169,3 +174,32 @@ class LibTest(TestCase):
                                         config_file=None,
                                         bot_config_file=None,
                                         bot_name='testbot')
+
+    def test_upload_file(self):
+        client, handler = self._create_client_and_handler_for_file_upload()
+        file = io.BytesIO(b'binary')
+
+        handler.upload_file(file)
+
+        client.upload_file.assert_called_once_with(file)
+
+    def test_upload_file_from_path(self):
+        client, handler = self._create_client_and_handler_for_file_upload()
+        file = io.BytesIO(b'binary')
+
+        with patch('builtins.open', return_value=file):
+            handler.upload_file_from_path('file.txt')
+
+        client.upload_file.assert_called_once_with(file)
+
+    def _create_client_and_handler_for_file_upload(self):
+        client = FakeClient()
+        client.upload_file = MagicMock()
+
+        handler = ExternalBotHandler(
+            client=client,
+            root_dir=None,
+            bot_details=None,
+            bot_config_file=None
+        )
+        return client, handler