congsh
44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
"""分类/标签模块测试"""
|
|
from app.tagger import classify_article, tag_article, tag_articles
|
|
from models import EnrichedArticle, Taxonomy
|
|
|
|
|
|
def test_classify_article(db):
|
|
categories = [
|
|
Taxonomy(name="科技", kind="category", keywords=["AI", "大模型", "芯片"]),
|
|
Taxonomy(name="财经", kind="category", keywords=["股市", "基金", "财报"]),
|
|
]
|
|
article = EnrichedArticle(
|
|
rk_article_id=1,
|
|
title="OpenAI 发布新一代大模型",
|
|
)
|
|
assert classify_article(article, categories) == "科技"
|
|
|
|
|
|
def test_tag_article(db):
|
|
tags = [
|
|
Taxonomy(name="人工智能", kind="tag", keywords=["AI", "人工智能", "大模型"]),
|
|
Taxonomy(name="半导体", kind="tag", keywords=["芯片", "半导体"]),
|
|
]
|
|
article = EnrichedArticle(
|
|
rk_article_id=1,
|
|
title="OpenAI 发布新一代大模型",
|
|
)
|
|
result = tag_article(article, tags)
|
|
assert "人工智能" in result
|
|
|
|
|
|
def test_tag_articles_integration(db):
|
|
db.add_all([
|
|
Taxonomy(name="科技", kind="category", keywords=["AI", "大模型"]),
|
|
Taxonomy(name="人工智能", kind="tag", keywords=["AI", "大模型"]),
|
|
])
|
|
article = EnrichedArticle(rk_article_id=1, title="OpenAI 发布 GPT-5 大模型")
|
|
db.add(article)
|
|
db.commit()
|
|
|
|
count = tag_articles(db)
|
|
assert count == 1
|
|
assert article.category == "科技"
|
|
assert "人工智能" in article.tags
|