feat: 添加 PostCrawlPipeline 统一管道 + 筛选服务 + 测试

This commit is contained in:
2026-05-09 13:57:37 +08:00
parent 022d348052
commit cad1390534
4 changed files with 286 additions and 0 deletions
+53
View File
@@ -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
+102
View File
@@ -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"
+93
View File
@@ -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()