132 lines
5.0 KiB
Python
132 lines
5.0 KiB
Python
import asyncio
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from db.database import init_db
|
|
from db.repository import Repository
|
|
from db.models import TargetChannel, SourceWebsite, Post
|
|
from services.website_analyzer import WebsiteAnalyzer
|
|
from services.admin_bot import get_persian_main_menu
|
|
|
|
|
|
async def test_target_context_message_count_and_history():
|
|
await init_db()
|
|
repo = Repository()
|
|
|
|
# 1. Create target
|
|
unique_channel_id = -10077665544 - int(time.time() % 100000)
|
|
target_id = await repo.add_target(channel_id=unique_channel_id, title="Target Context Test", username="trg_ctx_test")
|
|
assert target_id is not None
|
|
|
|
# Check default context count is 0
|
|
t = await repo.get_target_by_id(target_id)
|
|
assert t.context_message_count == 0
|
|
|
|
# Update context count to 5
|
|
await repo.update_target_context_count(target_id, 5)
|
|
t_updated = await repo.get_target_by_id(target_id)
|
|
assert t_updated.context_message_count == 5
|
|
|
|
# Insert published posts for this target
|
|
t0 = datetime(2026, 8, 28, 14, 0, 0, tzinfo=timezone.utc)
|
|
t1 = datetime(2026, 8, 28, 15, 0, 0, tzinfo=timezone.utc)
|
|
t2 = datetime(2026, 8, 28, 16, 0, 0, tzinfo=timezone.utc)
|
|
|
|
src_id = -10099000 - int(time.time() % 100000)
|
|
p0_id = await repo.create_raw_post(
|
|
source_channel_id=src_id,
|
|
source_message_id=501,
|
|
raw_text="پست منتشرشده ۱ در تارگت",
|
|
source_created_at=t0
|
|
)
|
|
await repo.record_post_published_to_target(p0_id, target_id, "Target Context Test")
|
|
|
|
p1_id = await repo.create_raw_post(
|
|
source_channel_id=src_id,
|
|
source_message_id=502,
|
|
raw_text="پست منتشرشده ۲ در تارگت",
|
|
source_created_at=t1
|
|
)
|
|
await repo.record_post_published_to_target(p1_id, target_id, "Target Context Test")
|
|
|
|
p2_id = await repo.create_raw_post(
|
|
source_channel_id=src_id,
|
|
source_message_id=503,
|
|
raw_text="پست در حال بازنویسی ۳",
|
|
source_created_at=t2
|
|
)
|
|
|
|
# Fetch recent target posts for target_id excluding p2_id -> should return p0 and p1 in chronological order
|
|
recent_target_posts = await repo.get_recent_target_posts(target_id, limit=5, exclude_post_id=p2_id)
|
|
assert len(recent_target_posts) == 2
|
|
assert recent_target_posts[0].id == p0_id
|
|
assert recent_target_posts[1].id == p1_id
|
|
|
|
# Clean up
|
|
await repo.delete_target(target_id)
|
|
|
|
|
|
async def test_website_custom_needs_and_prompt_injection():
|
|
await init_db()
|
|
repo = Repository()
|
|
|
|
test_url = f"https://techblog.example.com/site-{int(time.time())}"
|
|
site_id = await repo.add_source_website(
|
|
name="Tech Blog Custom Needs",
|
|
url=test_url,
|
|
custom_instructions="فقط اخبار مربوط به پردازندههای گرافیکی و تراشهها"
|
|
)
|
|
assert site_id is not None
|
|
|
|
site = await repo.get_source_website_by_id(site_id)
|
|
assert site.custom_instructions == "فقط اخبار مربوط به پردازندههای گرافیکی و تراشهها"
|
|
assert site.api_config == {} # Not analyzed automatically
|
|
|
|
# Test WebsiteAnalyzer prompt receives custom_instructions
|
|
captured_payload = {}
|
|
llm_mock = MagicMock()
|
|
async def fake_generate_json(prompt, system_prompt, action_name):
|
|
captured_payload["prompt"] = prompt
|
|
return {
|
|
"status": "success",
|
|
"endpoint_url": f"{test_url}/feed",
|
|
"parser_type": "rss",
|
|
"field_mappings": {"title": "title", "content": "description", "link": "link"}
|
|
}
|
|
llm_mock.generate_json = AsyncMock(side_effect=fake_generate_json)
|
|
|
|
analyzer = WebsiteAnalyzer(llm=llm_mock)
|
|
with patch("httpx.AsyncClient.get") as mock_get:
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.text = '<html><head><link rel="alternate" type="application/rss+xml" href="/feed" /></head></html>'
|
|
mock_get.return_value = mock_resp
|
|
|
|
ok, cfg, summary = await analyzer.analyze_website(test_url, custom_instructions=site.custom_instructions)
|
|
assert ok is True
|
|
assert "فقط اخبار مربوط به پردازندههای گرافیکی و تراشهها" in captured_payload["prompt"]
|
|
|
|
# Clean up
|
|
await repo.delete_source_website(site_id)
|
|
|
|
|
|
async def test_menu_layout():
|
|
menu = get_persian_main_menu(is_paused=False)
|
|
all_texts = [getattr(getattr(b, "button", None), "text", str(b)) for row in menu for b in row]
|
|
assert any("Copy" in t for t in all_texts)
|
|
assert any("AI" in t for t in all_texts)
|
|
assert any("Bots" in t for t in all_texts)
|
|
assert any("Monitor" in t for t in all_texts)
|
|
assert any("System" in t for t in all_texts)
|
|
|
|
|
|
async def main():
|
|
await test_target_context_message_count_and_history()
|
|
await test_website_custom_needs_and_prompt_injection()
|
|
await test_menu_layout()
|
|
print("All Target Context, Website Custom Needs, and Reorganized Menu tests passed successfully!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|