test: improve test assertions and edge cases

1 files changed, 27 insertions, 13 deletions

diff --git a/tests/cogs/test_ai.py b/tests/cogs/test_ai.py
index bd8f813..a1c3fd4 100644
--- a/tests/cogs/test_ai.py
+++ b/tests/cogs/test_ai.py
@@ -1,6 +1,7 @@
 import pytest
 from unittest.mock import AsyncMock, patch, MagicMock
 import discord
+import httpx
 from tests.conftest import MockApplicationContext
 
 
@@ -27,12 +28,18 @@ def mock_ai_response():
 @pytest.mark.asyncio
 async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot, mock_ai_response):
     ctx = MockApplicationContext(bot=bot, channel=mock_channel, guild=mock_guild)
-    
+
     mock_response = MagicMock()
     mock_response.json = MagicMock(return_value=mock_ai_response)
-    mock_response.status_code = 200
-    
-    with patch('cogs.ai.requests.post', return_value=mock_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.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
             await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
             mock_followup.assert_called_once_with("Test response from AI")
@@ -41,12 +48,18 @@ async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot, mo
 @pytest.mark.asyncio
 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_response.status_code = 200
-    
-    with patch('cogs.ai.requests.post', return_value=mock_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.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
             await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
             mock_followup.assert_called_once_with("bot siem zesroł")
@@ -55,15 +68,16 @@ async def test_senator_command_empty_choices(ai_cog, mock_channel, mock_guild, b
 @pytest.mark.asyncio
 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={})
-    mock_response.status_code = 200
-    
-    with patch('cogs.ai.requests.post', return_value=mock_response):
+
+    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
\ No newline at end of file
+            assert mock_logger.info.called