feat: target channel context history, custom website extraction needs, and categorized main menu

2026-08-28 22:14:49 +03:30
parent 8ccecfe699
commit 9551ccce71
8 changed files with 350 additions and 73 deletions
@@ -0,0 +1,129 @@
+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)
+
+    p0_id = await repo.create_raw_post(
+        source_channel_id=-1001,
+        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=-1001,
+        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=-1001,
+        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)
+    assert len(menu) == 5
+    # Row 1 has Home/Start
+    assert any("🏠 منوی اصلی" in btn.text for btn in menu[0])
+    # Row 4 has Unified Add
+    assert any("➕ افزودن مبدا / مقصد" in btn.text for btn in menu[3])
+
+
+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())