147 lines
6.7 KiB
Python
147 lines
6.7 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import apinex_client
|
|
import bot
|
|
import user_store
|
|
|
|
|
|
class FakeResponse:
|
|
def __init__(self, status, payload):
|
|
self.status = status
|
|
self.payload = payload
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
async def json(self, content_type=None):
|
|
return self.payload
|
|
|
|
|
|
class FakeClientSession:
|
|
def __init__(self, response, **kwargs):
|
|
self.response = response
|
|
self.kwargs = kwargs
|
|
self.request: tuple[str, dict] = ("", {})
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
def post(self, url, **kwargs):
|
|
self.request = (url, kwargs)
|
|
return self.response
|
|
|
|
|
|
class ApinexTests(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp_dir.cleanup)
|
|
self.db_path = str(Path(self.temp_dir.name) / "test.sqlite3")
|
|
self.db_env = patch.dict("os.environ", {"BOT_DB_PATH": self.db_path}, clear=False)
|
|
self.db_env.start()
|
|
self.addCleanup(self.db_env.stop)
|
|
user_store.init_db(self.db_path)
|
|
self.env = patch.dict("os.environ", {"APINEX_API_KEY": "not-logged"}, clear=False)
|
|
self.env.start()
|
|
self.addCleanup(self.env.stop)
|
|
|
|
def test_extract_transcript(self):
|
|
payload = {"results": {"channels": [{"alternatives": [{"transcript": " Привет, мир! "}]}]}}
|
|
self.assertEqual(apinex_client.extract_transcript(payload), "Привет, мир!")
|
|
self.assertEqual(apinex_client.extract_transcript({}), "")
|
|
|
|
async def test_transcribe_ogg_and_model_config(self):
|
|
response = FakeResponse(200, {"results": {"channels": [{"alternatives": [{"transcript": "Готово"}]}]}})
|
|
fake_session = FakeClientSession(response)
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=fake_session), patch.dict(
|
|
"os.environ", {"APINEX_STT_MODEL": "nova-2"}, clear=False
|
|
):
|
|
result = await apinex_client.transcribe_audio(b"ogg-data", "audio/ogg")
|
|
self.assertEqual(result, "Готово")
|
|
self.assertIsNotNone(fake_session.request)
|
|
url, kwargs = fake_session.request
|
|
self.assertEqual(url, apinex_client.APINEX_STT_URL)
|
|
self.assertEqual(kwargs["params"]["language"], "ru")
|
|
self.assertEqual(kwargs["params"]["model"], "nova-2")
|
|
self.assertEqual(kwargs["headers"]["Content-Type"], "audio/ogg")
|
|
self.assertEqual(kwargs["headers"]["Authorization"], "Bearer not-logged")
|
|
self.assertIn("Mozilla/5.0", kwargs["headers"]["User-Agent"])
|
|
|
|
async def test_chat_completion_uses_selected_model_and_parses_response(self):
|
|
response = FakeResponse(200, {"choices": [{"message": {"content": "Ответ модели"}}]})
|
|
fake_session = FakeClientSession(response)
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=fake_session), patch.dict(
|
|
"os.environ", {"APINEX_POST_WRITER_MODEL": "gemini-3.8-flash"}, clear=False
|
|
):
|
|
result = await apinex_client.chat_text("APINEX_POST_WRITER_MODEL", "system", "user")
|
|
self.assertEqual(result, "Ответ модели")
|
|
_, kwargs = fake_session.request
|
|
self.assertEqual(kwargs["json"]["model"], "gemini-3.8-flash")
|
|
self.assertEqual(kwargs["json"]["messages"][0]["role"], "system")
|
|
|
|
async def test_image_is_sent_as_private_inline_data_url(self):
|
|
response = FakeResponse(200, {"choices": [{"message": {"content": '{"posts":["Текст публикации"]}'}}]})
|
|
fake_session = FakeClientSession(response)
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=fake_session):
|
|
posts = await apinex_client.extract_posts_from_images([(b"PNG", "image/png")])
|
|
self.assertEqual(posts, ["Текст публикации"])
|
|
_, kwargs = fake_session.request
|
|
content = kwargs["json"]["messages"][1]["content"]
|
|
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/png;base64,"))
|
|
|
|
async def test_image_without_readable_posts_returns_empty_samples(self):
|
|
response = FakeResponse(200, {"choices": [{"message": {"content": '{"posts":[]}'}}]})
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=FakeClientSession(response)):
|
|
self.assertEqual(await apinex_client.extract_posts_from_images([(b"img", "image/jpeg")]), [])
|
|
|
|
async def test_intent_validation_and_three_posts(self):
|
|
with patch.object(apinex_client, "chat_text", return_value='{"intent":"write_post","topic":"кофе"}'):
|
|
intent = await apinex_client.classify_intent("придумай пост про кофе")
|
|
self.assertEqual(intent["intent"], "write_post")
|
|
with patch.object(apinex_client, "chat_text", return_value='{"posts":["A","B","C"]}'):
|
|
posts = await apinex_client.generate_posts(intent, "дружелюбный")
|
|
self.assertEqual(posts, ["A", "B", "C"])
|
|
with patch.object(apinex_client, "chat_text", return_value='{"posts":["A", "B"]}'):
|
|
with self.assertRaises(ValueError):
|
|
await apinex_client.generate_posts(intent, "")
|
|
|
|
async def test_http_errors_surface_without_secret_logging(self):
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=FakeClientSession(FakeResponse(401, {}))):
|
|
with self.assertRaisesRegex(RuntimeError, "HTTP 401"):
|
|
await apinex_client.transcribe_audio(b"ogg-data", "audio/ogg")
|
|
|
|
|
|
class StorageTests(unittest.TestCase):
|
|
def test_user_tone_isolated_per_id(self):
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db = str(Path(temp_dir) / "users.sqlite3")
|
|
user_store.init_db(db)
|
|
self.assertEqual(user_store.get_tone_of_voice(1, db), "")
|
|
user_store.save_tone_of_voice(1, "дружелюбный", db)
|
|
user_store.save_tone_of_voice(2, "официальный", db)
|
|
self.assertEqual(user_store.get_tone_of_voice(1, db), "дружелюбный")
|
|
self.assertEqual(user_store.get_tone_of_voice(2, db), "официальный")
|
|
|
|
|
|
class TelegramUtilsTests(unittest.TestCase):
|
|
def test_split_text_preserves_tokens_and_limit(self):
|
|
chunks = bot.split_text("слово " * 2000, 4000)
|
|
self.assertTrue(all(len(chunk) <= 4000 for chunk in chunks))
|
|
self.assertEqual(" ".join(chunks).split(), ("слово " * 2000).split())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|