48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
import sys
|
||
from pathlib import Path
|
||
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))
|
||
from models.episode import Episode
|
||
from distill.hard_rules import (
|
||
passes_rule_filter, content_density, is_small_talk,
|
||
is_knowledge_qa, has_user_message, has_assistant_message
|
||
)
|
||
|
||
def test_pass_empty_content():
|
||
ep = Episode.create(content='')
|
||
passed, reason = passes_rule_filter(ep)
|
||
assert not passed and '无用户消息' in reason
|
||
|
||
def test_pass_small_talk():
|
||
ep = Episode.create(content='hello', entities=[], facts=[])
|
||
passed, reason = passes_rule_filter(ep)
|
||
assert not passed and '闲聊' in reason
|
||
|
||
def test_pass_normal_content():
|
||
ep = Episode.create(content='我需要安装 Docker，请帮我配置')
|
||
passed, reason = passes_rule_filter(ep)
|
||
assert passed, reason
|
||
|
||
def test_content_density():
|
||
ep = Episode.create(content='a b c')
|
||
assert content_density(ep) == 0.0
|
||
ep2 = Episode.create(content='Python 是一种编程语言')
|
||
assert content_density(ep2) > 0
|
||
|
||
def test_small_talk():
|
||
ep = Episode.create(content='早安')
|
||
assert is_small_talk(ep)
|
||
|
||
def test_engine_stats_works():
|
||
from distill.engine import DistillEngine
|
||
eng = DistillEngine(':memory:')
|
||
stats = eng.run_batch(limit=0)
|
||
assert 'processed' in stats
|
||
|
||
if __name__ == '__main__':
|
||
test_pass_empty_content()
|
||
test_pass_small_talk()
|
||
test_pass_normal_content()
|
||
test_content_density()
|
||
test_small_talk()
|
||
test_engine_stats_works()
|
||
print('All hard_rules + engine basic tests passed!') |