183 lines
8.9 KiB
Python
183 lines
8.9 KiB
Python
import asyncio
|
|
import json
|
|
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):
|
|
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, "Ответ модели")
|
|
self.assertIsNotNone(fake_session.request)
|
|
_, kwargs = fake_session.request
|
|
self.assertEqual(kwargs["json"]["model"], "gemini-3.8-flash")
|
|
self.assertEqual(kwargs["json"]["messages"][0]["role"], "system")
|
|
|
|
async def test_single_image_is_sent_as_private_inline_data_url(self):
|
|
response = FakeResponse(200, {"choices": [{"message": {"content": "Дружелюбно и кратко"}}]})
|
|
fake_session = FakeClientSession(response)
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=fake_session):
|
|
tone = await apinex_client.extract_tone_from_image(b"PNG", "image/png")
|
|
self.assertEqual(tone, "Дружелюбно и кратко")
|
|
self.assertIsNotNone(fake_session.request)
|
|
_, 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_multiple_images_are_sent_in_one_vision_request(self):
|
|
response = FakeResponse(200, {"choices": [{"message": {"content": "Единый стиль"}}]})
|
|
fake_session = FakeClientSession(response)
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=fake_session):
|
|
tone = await apinex_client.extract_tone_from_images(
|
|
[(b"one", "image/png"), (b"two", "image/jpeg"), (b"three", "image/webp")]
|
|
)
|
|
self.assertEqual(tone, "Единый стиль")
|
|
self.assertIsNotNone(fake_session.request)
|
|
_, kwargs = fake_session.request
|
|
content = kwargs["json"]["messages"][1]["content"]
|
|
self.assertEqual(len(content), 4)
|
|
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/png;base64,"))
|
|
self.assertTrue(content[3]["image_url"]["url"].startswith("data:image/webp;base64,"))
|
|
|
|
async def test_writer_preserves_multiline_post_without_labels(self):
|
|
multiline_post = "Привет!\n\nВот новость ☕"
|
|
with patch.object(apinex_client, "chat_text", return_value=json.dumps({"posts": [multiline_post, "Вариант 2", "Вариант 3"]})):
|
|
posts = await apinex_client.generate_posts(
|
|
{"intent": "write_post", "topic": "кофейня", "platform": "", "requirements": ""}, ""
|
|
)
|
|
self.assertEqual(posts[0], multiline_post)
|
|
self.assertFalse(posts[0].startswith("1/3"))
|
|
|
|
async def test_image_rejects_invalid_screenshot_marker(self):
|
|
response = FakeResponse(200, {"choices": [{"message": {"content": "NOT_POST_SCREENSHOT"}}]})
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=FakeClientSession(response)):
|
|
self.assertEqual(await apinex_client.extract_tone_from_image(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_writer_receives_tone_and_reply_fields(self):
|
|
captured = {}
|
|
|
|
async def fake_chat(model_setting, system_prompt, user_content):
|
|
captured.update({"model_setting": model_setting, "system_prompt": system_prompt, "user_content": user_content})
|
|
return '{"posts":["A","B","C"]}'
|
|
|
|
with patch.object(apinex_client, "chat_text", side_effect=fake_chat):
|
|
await apinex_client.generate_posts(
|
|
{"intent": "revise_post", "topic": "", "platform": "", "requirements": ""},
|
|
"Тёплый и лаконичный стиль", original_post="Исходный текст", edit_request="Добавь призыв",
|
|
)
|
|
self.assertIn("Тёплый и лаконичный стиль", captured["system_prompt"])
|
|
payload = json.loads(captured["user_content"])
|
|
self.assertEqual(payload["original_post"], "Исходный текст")
|
|
self.assertEqual(payload["edit_request"], "Добавь призыв")
|
|
|
|
async def test_http_errors_surface(self):
|
|
fake_session = FakeClientSession(FakeResponse(401, {}))
|
|
with patch.object(apinex_client.aiohttp, "ClientSession", return_value=fake_session):
|
|
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()
|