feat: 添加 PostCrawlPipeline 统一管道 + 筛选服务 + 测试
This commit is contained in:
@@ -0,0 +1,53 @@
|
|||||||
|
from datetime import date
|
||||||
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
|
|
||||||
|
def filter_by_keywords(announcements: List[Dict[str, Any]],
|
||||||
|
keywords: List[str]) -> List[Dict[str, Any]]:
|
||||||
|
if not keywords:
|
||||||
|
return announcements
|
||||||
|
|
||||||
|
filtered = []
|
||||||
|
for ann in announcements:
|
||||||
|
search_text = f"{ann.get('title', '')} {ann.get('purchase_name', '')}"
|
||||||
|
if any(kw in search_text for kw in keywords):
|
||||||
|
ann["keyword_matched"] = True
|
||||||
|
filtered.append(ann)
|
||||||
|
|
||||||
|
return filtered
|
||||||
|
|
||||||
|
|
||||||
|
def filter_by_date(announcements: List[Dict[str, Any]],
|
||||||
|
start_date: Optional[date] = None,
|
||||||
|
end_date: Optional[date] = None) -> List[Dict[str, Any]]:
|
||||||
|
if not start_date and not end_date:
|
||||||
|
return announcements
|
||||||
|
|
||||||
|
filtered = []
|
||||||
|
for ann in announcements:
|
||||||
|
pub_date = ann.get("publish_date")
|
||||||
|
if not pub_date:
|
||||||
|
continue
|
||||||
|
if isinstance(pub_date, date):
|
||||||
|
pub_date = pub_date
|
||||||
|
else:
|
||||||
|
pub_date = pub_date.date() if hasattr(pub_date, "date") else pub_date
|
||||||
|
|
||||||
|
if start_date and pub_date < start_date:
|
||||||
|
continue
|
||||||
|
if end_date and pub_date > end_date:
|
||||||
|
continue
|
||||||
|
filtered.append(ann)
|
||||||
|
|
||||||
|
return filtered
|
||||||
|
|
||||||
|
|
||||||
|
def dedup_by_hash(announcements: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||||
|
seen = set()
|
||||||
|
result = []
|
||||||
|
for ann in announcements:
|
||||||
|
h = ann.get("content_hash")
|
||||||
|
if h and h not in seen:
|
||||||
|
seen.add(h)
|
||||||
|
result.append(ann)
|
||||||
|
return result
|
||||||
@@ -0,0 +1,102 @@
|
|||||||
|
from typing import Any, Dict, List
|
||||||
|
from app.crawler.base import PipelineConfig, PipelineResult
|
||||||
|
from app.services.filter_service import filter_by_keywords, dedup_by_hash
|
||||||
|
|
||||||
|
|
||||||
|
class PostCrawlPipeline:
|
||||||
|
def __init__(self, db_session, notification_service):
|
||||||
|
self.db = db_session
|
||||||
|
self.notify = notification_service
|
||||||
|
|
||||||
|
async def process(self, announcements: List[Dict[str, Any]],
|
||||||
|
config: PipelineConfig) -> PipelineResult:
|
||||||
|
result = PipelineResult()
|
||||||
|
|
||||||
|
if not announcements:
|
||||||
|
return result
|
||||||
|
|
||||||
|
# 1. Dedup
|
||||||
|
if config.dedup_enabled:
|
||||||
|
announcements = dedup_by_hash(announcements)
|
||||||
|
|
||||||
|
# 2. Store to database
|
||||||
|
stored = await self._save_to_db(announcements)
|
||||||
|
result.stored = stored
|
||||||
|
|
||||||
|
to_notify = announcements
|
||||||
|
|
||||||
|
# 3. Filter
|
||||||
|
if config.filter_enabled and config.keywords:
|
||||||
|
before = len(to_notify)
|
||||||
|
to_notify = [a for a in to_notify
|
||||||
|
if self._match_keywords(a, config.keywords)]
|
||||||
|
result.filtered = before - len(to_notify)
|
||||||
|
|
||||||
|
# 4. Notify
|
||||||
|
if config.notify_mode == "all":
|
||||||
|
result.notified = await self._send_notifications(to_notify)
|
||||||
|
elif config.notify_mode == "filtered":
|
||||||
|
if config.filter_enabled and config.keywords:
|
||||||
|
result.notified = await self._send_notifications(to_notify)
|
||||||
|
elif not config.filter_enabled:
|
||||||
|
result.notified = await self._send_notifications(to_notify)
|
||||||
|
|
||||||
|
# 5. Mark sent
|
||||||
|
if config.mark_sent and result.notified > 0:
|
||||||
|
await self._mark_sent(to_notify)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
async def _save_to_db(self, announcements: List[Dict[str, Any]]) -> int:
|
||||||
|
from sqlalchemy.dialects.postgresql import insert
|
||||||
|
from app.models.announcement import Announcement
|
||||||
|
|
||||||
|
if not announcements:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
values = [{
|
||||||
|
"title": a["title"],
|
||||||
|
"publish_date": a["publish_date"],
|
||||||
|
"purchase_name": a.get("purchase_name", ""),
|
||||||
|
"content_url": a.get("content_url", ""),
|
||||||
|
"source_code": a["source_code"],
|
||||||
|
"source_name": a["source_name"],
|
||||||
|
"announcement_type": a.get("announcement_type", "purchase"),
|
||||||
|
"content_hash": a["content_hash"],
|
||||||
|
"crawl_mode": a.get("crawl_mode", "auto"),
|
||||||
|
"is_new": a.get("is_new", True),
|
||||||
|
"is_sent": False,
|
||||||
|
"keyword_matched": a.get("keyword_matched", False),
|
||||||
|
} for a in announcements]
|
||||||
|
|
||||||
|
stmt = insert(Announcement).values(values)
|
||||||
|
stmt = stmt.on_conflict_do_nothing(index_elements=["content_hash"])
|
||||||
|
|
||||||
|
result_proxy = await self.db.execute(stmt)
|
||||||
|
await self.db.commit()
|
||||||
|
return result_proxy.rowcount or len(values)
|
||||||
|
|
||||||
|
async def _send_notifications(self, announcements: List[Dict[str, Any]]) -> int:
|
||||||
|
return await self.notify.send(announcements)
|
||||||
|
|
||||||
|
async def _mark_sent(self, announcements: List[Dict[str, Any]]) -> int:
|
||||||
|
from app.models.announcement import Announcement
|
||||||
|
from sqlalchemy import update
|
||||||
|
|
||||||
|
hashes = [a["content_hash"] for a in announcements if a.get("content_hash")]
|
||||||
|
if not hashes:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
stmt = (
|
||||||
|
update(Announcement)
|
||||||
|
.where(Announcement.content_hash.in_(hashes))
|
||||||
|
.values(is_sent=True)
|
||||||
|
)
|
||||||
|
result = await self.db.execute(stmt)
|
||||||
|
await self.db.commit()
|
||||||
|
return result.rowcount
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _match_keywords(announcement: Dict[str, Any], keywords: List[str]) -> bool:
|
||||||
|
text = f"{announcement.get('title', '')} {announcement.get('purchase_name', '')}"
|
||||||
|
return any(kw in text for kw in keywords)
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
from datetime import datetime
|
||||||
|
from app.services.filter_service import filter_by_keywords, filter_by_date
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_by_keywords_match():
|
||||||
|
announcements = [
|
||||||
|
{"title": "大化县采购公告", "purchase_name": "大化县财政局",
|
||||||
|
"source_code": "test", "source_name": "test"},
|
||||||
|
{"title": "南宁市采购公告", "purchase_name": "南宁市财政局",
|
||||||
|
"source_code": "test", "source_name": "test"},
|
||||||
|
]
|
||||||
|
result = filter_by_keywords(announcements, ["大化"])
|
||||||
|
assert len(result) == 1
|
||||||
|
assert result[0]["title"] == "大化县采购公告"
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_by_keywords_no_keywords():
|
||||||
|
announcements = [
|
||||||
|
{"title": "大化县采购公告", "purchase_name": "x",
|
||||||
|
"source_code": "test", "source_name": "test"},
|
||||||
|
]
|
||||||
|
result = filter_by_keywords(announcements, [])
|
||||||
|
assert len(result) == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_by_date_range():
|
||||||
|
today = datetime(2026, 5, 9)
|
||||||
|
announcements = [
|
||||||
|
{"title": "t1", "publish_date": datetime(2026, 5, 9),
|
||||||
|
"source_code": "test", "source_name": "test"},
|
||||||
|
{"title": "t2", "publish_date": datetime(2026, 5, 1),
|
||||||
|
"source_code": "test", "source_name": "test"},
|
||||||
|
{"title": "t3", "publish_date": datetime(2026, 4, 30),
|
||||||
|
"source_code": "test", "source_name": "test"},
|
||||||
|
]
|
||||||
|
result = filter_by_date(announcements, start_date=today, end_date=today)
|
||||||
|
assert len(result) == 1
|
||||||
|
assert result[0]["title"] == "t1"
|
||||||
@@ -0,0 +1,93 @@
|
|||||||
|
import pytest
|
||||||
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
from datetime import datetime
|
||||||
|
from app.services.pipeline import PostCrawlPipeline
|
||||||
|
from app.crawler.base import CrawlResult, PipelineConfig
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_pipeline_filtered_mode():
|
||||||
|
config = PipelineConfig(
|
||||||
|
filter_enabled=True,
|
||||||
|
keywords=["大化"],
|
||||||
|
dedup_enabled=True,
|
||||||
|
notify_mode="filtered",
|
||||||
|
mark_sent=False,
|
||||||
|
)
|
||||||
|
result = CrawlResult(
|
||||||
|
source_code="test", source_name="test",
|
||||||
|
total_count=3, new_count=3,
|
||||||
|
announcements=[
|
||||||
|
{
|
||||||
|
"title": "大化县公告", "publish_date": datetime(2026, 5, 9),
|
||||||
|
"purchase_name": "大化县", "content_url": "https://1.com",
|
||||||
|
"source_code": "test", "source_name": "测试",
|
||||||
|
"announcement_type": "purchase", "crawl_mode": "auto",
|
||||||
|
"is_new": True, "is_today": True,
|
||||||
|
"content_hash": "abc123",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"title": "南宁市公告", "publish_date": datetime(2026, 5, 9),
|
||||||
|
"purchase_name": "南宁市", "content_url": "https://2.com",
|
||||||
|
"source_code": "test", "source_name": "测试",
|
||||||
|
"announcement_type": "purchase", "crawl_mode": "auto",
|
||||||
|
"is_new": True, "is_today": True,
|
||||||
|
"content_hash": "def456",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_db = AsyncMock()
|
||||||
|
mock_notify = AsyncMock()
|
||||||
|
|
||||||
|
pipeline = PostCrawlPipeline(db_session=mock_db, notification_service=mock_notify)
|
||||||
|
|
||||||
|
with patch.object(pipeline, "_save_to_db", AsyncMock(return_value=2)):
|
||||||
|
with patch.object(pipeline, "_send_notifications", AsyncMock(return_value=1)):
|
||||||
|
pipe_result = await pipeline.process(
|
||||||
|
result.announcements, config
|
||||||
|
)
|
||||||
|
assert pipe_result.stored == 2
|
||||||
|
assert pipe_result.filtered == 1
|
||||||
|
assert pipe_result.notified == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_pipeline_all_mode():
|
||||||
|
config = PipelineConfig(
|
||||||
|
filter_enabled=False,
|
||||||
|
keywords=[],
|
||||||
|
dedup_enabled=True,
|
||||||
|
notify_mode="all",
|
||||||
|
mark_sent=True,
|
||||||
|
)
|
||||||
|
result = CrawlResult(
|
||||||
|
source_code="dahuagov", source_name="大化县政府网",
|
||||||
|
total_count=2, new_count=2,
|
||||||
|
announcements=[
|
||||||
|
{
|
||||||
|
"title": f"公告{i}", "publish_date": datetime(2026, 5, 9),
|
||||||
|
"purchase_name": "大化县", "content_url": f"https://x.com/{i}",
|
||||||
|
"source_code": "dahuagov", "source_name": "大化县政府网采购公告",
|
||||||
|
"announcement_type": "purchase", "crawl_mode": "auto",
|
||||||
|
"is_new": True, "is_today": True,
|
||||||
|
"content_hash": f"hash{i}",
|
||||||
|
}
|
||||||
|
for i in range(2)
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_db = AsyncMock()
|
||||||
|
mock_notify = AsyncMock()
|
||||||
|
pipeline = PostCrawlPipeline(db_session=mock_db, notification_service=mock_notify)
|
||||||
|
|
||||||
|
with patch.object(pipeline, "_save_to_db", AsyncMock(return_value=2)):
|
||||||
|
with patch.object(pipeline, "_send_notifications", AsyncMock(return_value=2)):
|
||||||
|
with patch.object(pipeline, "_mark_sent", AsyncMock(return_value=2)):
|
||||||
|
pipe_result = await pipeline.process(
|
||||||
|
result.announcements, config
|
||||||
|
)
|
||||||
|
assert pipe_result.stored == 2
|
||||||
|
assert pipe_result.filtered == 0
|
||||||
|
assert pipe_result.notified == 2
|
||||||
|
pipeline._mark_sent.assert_awaited_once()
|
||||||
Reference in New Issue
Block a user