111 lines
4.3 KiB
Python
111 lines
4.3 KiB
Python
import asyncio
|
||
import time
|
||
from datetime import datetime, timezone, timedelta
|
||
from unittest.mock import AsyncMock, patch, MagicMock
|
||
from db.database import init_db
|
||
from db.repository import Repository
|
||
from db.models import TargetChannel, Post, AIProviderProfile
|
||
from services.ai_processor import AIProcessor
|
||
from core.llm import LLMClient
|
||
|
||
|
||
async def test_source_created_at_and_context_count():
|
||
await init_db()
|
||
repo = Repository()
|
||
|
||
unique_channel_id = -10099887700 - int(time.time() % 100000)
|
||
src_id = await repo.add_source(unique_channel_id, "Context Test Source", "ctx_test")
|
||
assert src_id is not None
|
||
|
||
# Check default context_message_count is 0
|
||
src = await repo.get_source_by_id(src_id)
|
||
assert src.context_message_count == 0
|
||
|
||
# Update context_message_count to 5
|
||
await repo.update_source_context_count(src_id, 5)
|
||
src_updated = await repo.get_source_by_id(src_id)
|
||
assert src_updated.context_message_count == 5
|
||
|
||
# Insert posts with explicit source_created_at timestamps
|
||
t0 = datetime(2026, 8, 28, 10, 0, 0, tzinfo=timezone.utc)
|
||
t1 = datetime(2026, 8, 28, 11, 0, 0, tzinfo=timezone.utc)
|
||
t2 = datetime(2026, 8, 28, 12, 0, 0, tzinfo=timezone.utc)
|
||
|
||
p0_id = await repo.create_raw_post(
|
||
source_channel_id=unique_channel_id,
|
||
source_message_id=101,
|
||
raw_text="خبر اول: مذاکرات آغاز شد.",
|
||
source_created_at=t0
|
||
)
|
||
p1_id = await repo.create_raw_post(
|
||
source_channel_id=unique_channel_id,
|
||
source_message_id=102,
|
||
raw_text="خبر دوم: توافقات اولیه حاصل گردید.",
|
||
source_created_at=t1
|
||
)
|
||
p2_id = await repo.create_raw_post(
|
||
source_channel_id=unique_channel_id,
|
||
source_message_id=103,
|
||
raw_text="خبر سوم: بیانیه مشترک امضا شد.",
|
||
source_created_at=t2
|
||
)
|
||
|
||
assert p0_id is not None
|
||
assert p1_id is not None
|
||
assert p2_id is not None
|
||
|
||
# Check source_created_at persisted
|
||
p2_loaded = await repo.get_post_by_id(p2_id)
|
||
assert p2_loaded.source_created_at is not None
|
||
|
||
# Fetch recent 2 posts excluding p2_id -> should return p0 and p1 in chronological order
|
||
recent_posts = await repo.get_recent_source_posts(unique_channel_id, limit=2, exclude_post_id=p2_id)
|
||
assert len(recent_posts) == 2
|
||
assert recent_posts[0].id == p0_id
|
||
assert recent_posts[1].id == p1_id
|
||
|
||
# Clean up
|
||
await repo.delete_source(src_id)
|
||
|
||
|
||
async def test_ai_processor_context_injection():
|
||
repo_mock = AsyncMock()
|
||
llm_mock = MagicMock()
|
||
captured_payload = {}
|
||
|
||
async def fake_generate_json(prompt, system_prompt, action_name, image_path=None):
|
||
captured_payload["prompt"] = prompt
|
||
captured_payload["system_prompt"] = system_prompt
|
||
return {"decision": "accept", "rewritten_text": "پست بازنویسیشده با در نظر گرفتن پیوستگی زمینه"}
|
||
|
||
llm_mock.generate_json = AsyncMock(side_effect=fake_generate_json)
|
||
|
||
processor = AIProcessor(repo=repo_mock, llm=llm_mock, double_check=False)
|
||
target = TargetChannel(id=1, channel_id=-1001234, title="Target News", username="trg_news", personality="رسمی و خبری")
|
||
|
||
ctx_p1 = Post(id=10, source_channel_id=-100, source_message_id=1, raw_text="پست زمینه ۱: مرحله اول آغاز شد.", source_created_at="2026-08-28 10:00:00")
|
||
ctx_p2 = Post(id=11, source_channel_id=-100, source_message_id=2, raw_text="پست زمینه ۲: مرحله دوم با موفقیت انجام شد.", source_created_at="2026-08-28 11:00:00")
|
||
|
||
result = await processor.rewrite_for_target(
|
||
raw_text="پست جدید ۳: نتایج نهایی اعلام شد.",
|
||
target=target,
|
||
context_posts=[ctx_p1, ctx_p2]
|
||
)
|
||
|
||
assert result.is_rejected is False
|
||
prompt_sent = captured_payload["prompt"]
|
||
assert "پیامها و پستهای قبلی/اخیر این کانال" in prompt_sent
|
||
assert "پست زمینه ۱" in prompt_sent
|
||
assert "پست زمینه ۲" in prompt_sent
|
||
assert "پست جدید ۳: نتایج نهایی اعلام شد." in prompt_sent
|
||
|
||
|
||
async def main():
|
||
await test_source_created_at_and_context_count()
|
||
await test_ai_processor_context_injection()
|
||
print("All source context history and original timestamp tests passed successfully!")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|