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 = '' 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) assert len(menu) == 5 first_row_texts = [getattr(getattr(b, "button", None), "text", str(b)) for b in menu[0]] assert any("🏠 منوی اصلی" in t for t in first_row_texts) fourth_row_texts = [getattr(getattr(b, "button", None), "text", str(b)) for b in menu[3]] assert any("➕ افزودن مبدا / مقصد" in t for t in fourth_row_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())