Support drawing.

tests/test_main.py

@@ -1,6 +1,10 @@
 import unittest
+import pytest
+import aiohttp
 import json
-from unittest.mock import Mock, PropertyMock, MagicMock, AsyncMock, patch, mock_open
+import openai
+import logging
+from unittest.mock import Mock, PropertyMock, MagicMock, AsyncMock, patch, mock_open, ANY
 from fjerkroa_bot import FjerkroaBot
 from discord import User, Message, TextChannel
 
@@ -14,10 +18,14 @@ class TestBotBase(unittest.IsolatedAsyncioTestCase):
         ]
         self.config_data = {
             "openai_key": "OPENAIKEY",
-            "engine": "gpt-4",
+            "model": "gpt-4",
             "max_tokens": 1024,
-            "n": 1,
             "temperature": 0.9,
+            "top-p": 1.0,
+            "presence-penalty": 1.0,
+            "frequency-penalty": 1.0,
+            "history-limit": 10,
+            "system": "You are an smart AI",
         }
         self.history_data = []
         with patch.object(FjerkroaBot, 'load_config', new=lambda s, c: self.config_data), \
@@ -25,11 +33,16 @@ class TestBotBase(unittest.IsolatedAsyncioTestCase):
             mock_user.return_value = MagicMock(spec=User)
             mock_user.return_value.id = 12
             self.bot = FjerkroaBot('config.json')
+            self.bot.staff_channel = AsyncMock(spec=TextChannel)
+            self.bot.staff_channel.send = AsyncMock()
+            self.bot.welcome_channel = AsyncMock(spec=TextChannel)
+            self.bot.welcome_channel.send = AsyncMock()
 
     def create_message(self, message: str) -> Message:
         message = MagicMock(spec=Message)
         message.content = "Hello, how are you?"
         message.author = AsyncMock(spec=User)
+        message.author.name = 'Lala'
         message.author.id = 123
         message.author.bot = False
         message.channel = AsyncMock(spec=TextChannel)
@@ -45,10 +58,78 @@ class TestFunctionality(TestBotBase):
         self.assertEqual(result, self.config_data)
 
     async def test_on_message_event(self) -> None:
+        async def acreate(*a, **kw):
+            answer = {'answer': 'Hello!',
+                      'answer_needed': True,
+                      'staff': None,
+                      'picture': None,
+                      'hack': False}
+            return {'choices': [{'message': {'content': json.dumps(answer)}}]}
         message = self.create_message("Hello there! How are you?")
-        await self.bot.on_message(message)
+        with patch.object(openai.ChatCompletion, 'acreate', new=acreate):
+            await self.bot.on_message(message)
         message.channel.send.assert_called_once_with("Hello!")
 
+    async def test_on_message_event2(self) -> None:
+        async def acreate(*a, **kw):
+            answer = {'answer': 'Hello!',
+                      'answer_needed': True,
+                      'staff': 'Hallo staff',
+                      'picture': None,
+                      'hack': False}
+            return {'choices': [{'message': {'content': json.dumps(answer)}}]}
+        message = self.create_message("Hello there! How are you?")
+        with patch.object(openai.ChatCompletion, 'acreate', new=acreate):
+            await self.bot.on_message(message)
+        message.channel.send.assert_called_once_with("Hello!")
+
+    async def test_on_message_event3(self) -> None:
+        async def acreate(*a, **kw):
+            return {'choices': [{'message': {'content': '{ "test": 3 ]'}}]}
+
+        def logging_warning(msg):
+            raise RuntimeError(msg)
+        message = self.create_message("Hello there! How are you?")
+        with patch.object(openai.ChatCompletion, 'acreate', new=acreate), \
+             patch.object(logging, 'warning', logging_warning):
+            with pytest.raises(RuntimeError, match='failed.*JSONDecodeError.*'):
+                await self.bot.on_message(message)
+
+    async def test_on_message_event4(self) -> None:
+        async def acreate(*a, **kw):
+            answer = {'answer': 'Hello!',
+                      'answer_needed': True,
+                      'staff': 'none',
+                      'picture': 'Some picture',
+                      'hack': False}
+            return {'choices': [{'message': {'content': json.dumps(answer)}}]}
+
+        async def adraw(*a, **kw):
+            return {'data': [{'url': 'http:url'}]}
+
+        def logging_warning(msg):
+            raise RuntimeError(msg)
+
+        class image:
+            def __init__(self, *args, **kw):
+                pass
+
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, exc_type, exc_value, traceback):
+                return False
+
+            async def read(self):
+                return b'test bytes'
+        message = self.create_message("Hello there! How are you?")
+        with patch.object(openai.ChatCompletion, 'acreate', new=acreate), \
+             patch.object(openai.Image, 'acreate', new=adraw), \
+             patch.object(logging, 'warning', logging_warning), \
+             patch.object(aiohttp.ClientSession, 'get', new=image):
+            await self.bot.on_message(message)
+        message.channel.send.assert_called_once_with("Hello!", files=[ANY])
+
 
 if __name__ == "__mait__":
     unittest.main()