1 files changed, 10 insertions, 50 deletions

diff --git a/tests/cogs/test_ai.py b/tests/cogs/test_ai.py
index 6d92f55..aab9861 100644
--- a/tests/cogs/test_ai.py
+++ b/tests/cogs/test_ai.py
@@ -1,7 +1,5 @@
 import pytest
 from unittest.mock import AsyncMock, patch, MagicMock
-import discord
-import httpx
 from conftest import MockApplicationContext
 
 
@@ -12,36 +10,14 @@ def ai_cog(bot):
     return cog
 
 
-@pytest.fixture
-def mock_ai_response():
-    return {
-        "choices": [
-            {
-                "message": {
-                    "content": "Test response from AI"
-                }
-            }
-        ]
-    }
-
-
 @pytest.mark.asyncio
-async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot, mock_ai_response):
+async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot):
     ctx = MockApplicationContext(bot=bot, channel=mock_channel, guild=mock_guild)
 
-    mock_response = MagicMock()
-    mock_response.json = MagicMock(return_value=mock_ai_response)
-
-    mock_client = MagicMock()
-    mock_client.post = AsyncMock(return_value=mock_response)
-    mock_client.__aenter__ = AsyncMock(return_value=mock_response)
-    mock_client.__aexit__ = AsyncMock(return_value=None)
-
-    with patch('cogs.ai.httpx.AsyncClient') as mock_httpx:
-        instance = mock_httpx.return_value.__aenter__.return_value
-        instance.post = AsyncMock(return_value=mock_response)
+    with patch('cogs.ai.llm_completion', new_callable=AsyncMock, return_value="Test response from AI") as mock_llm:
         with patch.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
             await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
+            mock_llm.assert_awaited_once()
             mock_followup.assert_called_once_with("Test response from AI")
 
 
@@ -49,19 +25,10 @@ async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot, mo
 async def test_senator_command_empty_choices(ai_cog, mock_channel, mock_guild, bot):
     ctx = MockApplicationContext(bot=bot, channel=mock_channel, guild=mock_guild)
 
-    mock_response = MagicMock()
-    mock_response.json = MagicMock(return_value={"choices": []})
-
-    mock_client = MagicMock()
-    mock_client.post = AsyncMock(return_value=mock_response)
-    mock_client.__aenter__ = AsyncMock(return_value=mock_response)
-    mock_client.__aexit__ = AsyncMock(return_value=None)
-
-    with patch('cogs.ai.httpx.AsyncClient') as mock_httpx:
-        instance = mock_httpx.return_value.__aenter__.return_value
-        instance.post = AsyncMock(return_value=mock_response)
+    with patch('cogs.ai.llm_completion', new_callable=AsyncMock, return_value=None) as mock_llm:
         with patch.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
             await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
+            mock_llm.assert_awaited_once()
             mock_followup.assert_called_once_with("bot siem zesroł")
 
 
@@ -69,15 +36,8 @@ async def test_senator_command_empty_choices(ai_cog, mock_channel, mock_guild, b
 async def test_senator_command_no_choices_key(ai_cog, mock_channel, mock_guild, bot):
     ctx = MockApplicationContext(bot=bot, channel=mock_channel, guild=mock_guild)
 
-    mock_response = MagicMock()
-    mock_response.json = MagicMock(return_value={})
-
-    with patch('cogs.ai.httpx.AsyncClient') as mock_httpx:
-        instance = mock_httpx.return_value.__aenter__.return_value
-        instance.post = AsyncMock(return_value=mock_response)
-        with patch('cogs.ai.logger') as mock_logger:
-            try:
-                await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
-            except KeyError:
-                pass
-            assert mock_logger.info.called
+    with patch('cogs.ai.llm_completion', new_callable=AsyncMock, return_value=None) as mock_llm:
+        with patch.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
+            await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
+            mock_llm.assert_awaited_once()
+            mock_followup.assert_called_once_with("bot siem zesroł")