import os import unittest from unittest.mock import AsyncMock, MagicMock, Mock, PropertyMock, mock_open, patch import toml from discord import Message, TextChannel, User from fjerkroa_bot import FjerkroaBot from fjerkroa_bot.ai_responder import AIMessage, AIResponse, parse_maybe_json 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.get("OPENAI_TOKEN", "test"), "model": "gpt-3.5-turbo", "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) json_struct = '{"This is a string."}' expected_output = "This is a string." self.assertEqual(parse_maybe_json(json_struct), expected_output) json_struct = '["This is a string."]' expected_output = "This is a string." self.assertEqual(parse_maybe_json(json_struct), expected_output) json_struct = "{This is a string.}" expected_output = "This is a string." self.assertEqual(parse_maybe_json(json_struct), expected_output) json_struct = "[This is a string.]" expected_output = "This is a string." self.assertEqual(parse_maybe_json(json_struct), expected_output) async def test_message_lings(self) -> None: request = AIMessage( "Lala", "Hello there!", "chat", False, ) message = { "answer": "Test [Link](https://www.example.com/test)", "answer_needed": True, "channel": "chat", "staff": None, "picture": None, "hack": False, } expected = AIResponse("Test https://www.example.com/test", True, "chat", None, None, False, False) self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected)) message = { "answer": "Test @[Link](https://www.example.com/test)", "answer_needed": True, "channel": "chat", "staff": None, "picture": None, "hack": False, } expected = AIResponse("Test Link", True, "chat", None, None, False, False) self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected)) message = { "answer": "Test [Link](https://www.example.com/test) and [Link2](https://xxx) lala", "answer_needed": True, "channel": "chat", "staff": None, "picture": None, "hack": False, } expected = AIResponse("Test https://www.example.com/test and https://xxx lala", True, "chat", None, None, False, False) self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected)) async def test_on_message_stort_path(self) -> None: message = self.create_message("Hello there! How are you?") message.author.name = "madeup_name" message.channel.name = "some_channel" # type: ignore self.bot.config["short-path"] = [[r"some.*", r"madeup.*"]] 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, "historise_question": true}', ) @patch("builtins.open", new_callable=mock_open) def test_update_history_with_file(self, mock_file): self.bot.airesponder.update_history({"content": '{"q": "What\'s your name?"}'}, {"content": '{"a": "AI"}'}, 10) self.assertEqual(len(self.bot.airesponder.history), 2) self.bot.airesponder.update_history({"content": '{"q1": "Q1"}'}, {"content": '{"a1": "A1"}'}, 2) self.bot.airesponder.update_history({"content": '{"q2": "Q2"}'}, {"content": '{"a2": "A2"}'}, 2) self.assertEqual(len(self.bot.airesponder.history), 2) self.bot.airesponder.history_file = "mock_file.pkl" self.bot.airesponder.update_history({"content": '{"q": "What\'s your favorite color?"}'}, {"content": '{"a": "Blue"}'}, 10) mock_file.assert_called_once_with("mock_file.pkl", "wb") mock_file().write.assert_called_once() if __name__ == "__mait__": unittest.main()