173 lines
7.3 KiB
Python
173 lines
7.3 KiB
Python
import os
|
|
import unittest
|
|
import aiohttp
|
|
import json
|
|
import toml
|
|
import openai
|
|
import logging
|
|
import pytest
|
|
from unittest.mock import Mock, PropertyMock, MagicMock, AsyncMock, patch, mock_open, ANY
|
|
from fjerkroa_bot import FjerkroaBot
|
|
from fjerkroa_bot.ai_responder import parse_maybe_json
|
|
from discord import User, Message, TextChannel
|
|
|
|
|
|
class TestBotBase(unittest.IsolatedAsyncioTestCase):
|
|
|
|
async def asyncSetUp(self):
|
|
self.mock_response = Mock()
|
|
self.mock_response.choices = [
|
|
Mock(text="Nice day today!")
|
|
]
|
|
self.config_data = {
|
|
"openai-token": os.environ['OPENAI_TOKEN'],
|
|
"model": "gpt-4",
|
|
"max-tokens": 1024,
|
|
"temperature": 0.9,
|
|
"top-p": 1.0,
|
|
"presence-penalty": 1.0,
|
|
"frequency-penalty": 1.0,
|
|
"history-limit": 10,
|
|
"system": "You are an smart AI",
|
|
"additional-responders": [],
|
|
}
|
|
self.history_data = []
|
|
with patch.object(FjerkroaBot, 'load_config', new=lambda s, c: self.config_data), \
|
|
patch.object(FjerkroaBot, 'user', new_callable=PropertyMock) as mock_user:
|
|
mock_user.return_value = MagicMock(spec=User)
|
|
mock_user.return_value.id = 12
|
|
self.bot = FjerkroaBot('config.toml')
|
|
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()
|
|
self.bot.airesponder.config = self.config_data
|
|
|
|
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)
|
|
message.channel.send = AsyncMock()
|
|
return message
|
|
|
|
|
|
class TestFunctionality(TestBotBase):
|
|
|
|
def test_load_config(self) -> None:
|
|
with patch('builtins.open', mock_open(read_data=toml.dumps(self.config_data))):
|
|
result = FjerkroaBot.load_config('config.toml')
|
|
self.assertEqual(result, self.config_data)
|
|
|
|
def test_json_strings(self) -> None:
|
|
json_string = '{"key1": "value1", "key2": "value2"}'
|
|
expected_output = "value1\nvalue2"
|
|
self.assertEqual(parse_maybe_json(json_string), expected_output)
|
|
non_json_string = "This is not a JSON string."
|
|
self.assertEqual(parse_maybe_json(non_json_string), non_json_string)
|
|
json_array = '["value1", "value2", "value3"]'
|
|
|
|
expected_output = "value1\nvalue2\nvalue3"
|
|
self.assertEqual(parse_maybe_json(json_array), expected_output)
|
|
json_string = '"value1"'
|
|
expected_output = 'value1'
|
|
self.assertEqual(parse_maybe_json(json_string), expected_output)
|
|
|
|
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?")
|
|
with patch.object(openai.ChatCompletion, 'acreate', new=acreate):
|
|
await self.bot.on_message(message)
|
|
message.channel.send.assert_called_once_with("Hello!", suppress_embeds=True)
|
|
|
|
async def test_on_message_stort_path(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?")
|
|
message.author.name = 'madeup_name'
|
|
message.channel.name = 'some_channel'
|
|
self.bot.config['short-path'] = [[r'some.*', r'madeup.*']]
|
|
with patch.object(openai.ChatCompletion, 'acreate', new=acreate):
|
|
await self.bot.on_message(message)
|
|
self.assertEqual(self.bot.airesponder.history[-1]["content"],
|
|
'{"user": "madeup_name", "message": "Hello, how are you?", "channel": "some_channel", "direct": false}')
|
|
message.author.name = 'different_name'
|
|
await self.bot.on_message(message)
|
|
self.assertEqual(self.bot.airesponder.history[-2]["content"],
|
|
'{"user": "different_name", "message": "Hello, how are you?", "channel": "some_channel", "direct": false}')
|
|
message.channel.send.assert_called_once_with("Hello!", suppress_embeds=True)
|
|
|
|
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!", suppress_embeds=True)
|
|
|
|
async def test_on_message_event3(self) -> None:
|
|
async def acreate(*a, **kw):
|
|
return {'choices': [{'message': {'content': '{ "test": 3 ]'}}]}
|
|
message = self.create_message("Hello there! How are you?")
|
|
with patch.object(openai.ChatCompletion, 'acreate', new=acreate):
|
|
with pytest.raises(RuntimeError, match="Failed to generate answer after multiple retries"):
|
|
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], suppress_embeds=True)
|
|
|
|
|
|
if __name__ == "__mait__":
|
|
unittest.main()
|