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