308 lines
15 KiB
Python
308 lines
15 KiB
Python
|
"""Background assessments use real SQLite and a local neutral knowledge index.
|
||
|
|
|
||
|
|
Only model generation/embedding are faked. Quick retrieval is run in process so
|
||
|
|
it uses the test's isolated index instead of starting the application's worker.
|
||
|
|
"""
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from fastapi import HTTPException
|
||
|
|
from pydantic import ValidationError
|
||
|
|
from sqlmodel import SQLModel, Session, create_engine
|
||
|
|
|
||
|
|
from app import config, training_map as tm
|
||
|
|
from app.models import AppUser, KnowledgeBase, KnowledgeBaseGrant, KnowledgeFile, ModelProvider
|
||
|
|
from app.rag import embed, pipeline, store
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def assessment_env(tmp_path, monkeypatch):
|
||
|
|
monkeypatch.setattr(config, 'DATA_DIR', tmp_path)
|
||
|
|
fallback = tmp_path / 'vectors'
|
||
|
|
fallback.mkdir()
|
||
|
|
monkeypatch.setattr(store, 'FALLBACK_DIR', fallback)
|
||
|
|
# The separate RAG process cannot inherit monkeypatched filesystem paths.
|
||
|
|
monkeypatch.setattr(tm, 'search_kbs_remote', pipeline.search_kbs)
|
||
|
|
monkeypatch.setattr(embed, 'embed_query', lambda text: [1.0, 0.0])
|
||
|
|
monkeypatch.setattr(tm, '_assessment_slots', threading.BoundedSemaphore(2))
|
||
|
|
monkeypatch.setattr(tm, '_assessment_jobs', set())
|
||
|
|
try:
|
||
|
|
from app import training_corpus
|
||
|
|
monkeypatch.setattr(training_corpus, 'CORPUS_DIR', tmp_path / 'corpus')
|
||
|
|
except ImportError:
|
||
|
|
pass
|
||
|
|
engine = create_engine(f'sqlite:///{tmp_path / "test.sqlite"}',
|
||
|
|
connect_args={'check_same_thread': False, 'timeout': 2})
|
||
|
|
SQLModel.metadata.create_all(engine)
|
||
|
|
folder = tmp_path / 'knowledge'
|
||
|
|
folder.mkdir()
|
||
|
|
text = '连续三天记录任务时间与打断来源，再按任务类别复盘时间分配并选择一个改进动作。'
|
||
|
|
(folder / 'lesson.txt').write_text(text, encoding='utf-8')
|
||
|
|
with Session(engine) as session:
|
||
|
|
user = AppUser(username='assessment-reader', role='editor')
|
||
|
|
kb = KnowledgeBase(name='中性时间管理资料', owner_user_id=99, visibility='shared',
|
||
|
|
source_dir=str(folder), file_count=1, chunk_count=1, index_status='ready')
|
||
|
|
provider = ModelProvider(name='test-model', default_model='neutral', api_key='unit-test-only',
|
||
|
|
base_url='http://127.0.0.1:9/v1', enabled=True)
|
||
|
|
session.add_all([user, kb, provider]); session.commit()
|
||
|
|
session.refresh(user); session.refresh(kb); session.refresh(provider)
|
||
|
|
grant = KnowledgeBaseGrant(kb_id=kb.id, user_id=user.id, permission='use')
|
||
|
|
session.add_all([grant, KnowledgeFile(kb_id=kb.id, file_name='lesson.txt')])
|
||
|
|
session.commit(); session.refresh(grant)
|
||
|
|
request = SimpleNamespace(state=SimpleNamespace(user=user))
|
||
|
|
row = tm.create_draft(tm.DraftInput(topic='时间管理训练营', kb_ids=[kb.id]), request, session)
|
||
|
|
ids = (row.id, kb.id, user.id, grant.id, provider.id)
|
||
|
|
store._write_fallback(kb.id, [{'id': 'time-1', 'text': text,
|
||
|
|
'meta': {'file': 'lesson.txt'}, 'embedding': [1.0, 0.0]}])
|
||
|
|
monkeypatch.setattr(tm, '_pick_model', lambda session: (ids[4], 'neutral'))
|
||
|
|
yield SimpleNamespace(engine=engine, draft_id=ids[0], kb_id=ids[1], user_id=ids[2], grant_id=ids[3])
|
||
|
|
engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
def request_for(session, env):
|
||
|
|
return SimpleNamespace(state=SimpleNamespace(user=session.get(AppUser, env.user_id)))
|
||
|
|
|
||
|
|
|
||
|
|
def neutral_plan(evidence):
|
||
|
|
return {'recommended_days': 1, 'rationale': '资料支持一天的记录和复盘练习。', 'insufficiencies': [],
|
||
|
|
'days': [{'day': 1, 'title': '记录时间分配', 'source_ids': [evidence[0]['id']],
|
||
|
|
'support': '存在具体记录步骤', 'prerequisite': '从零开始', 'outcome': '完成记录并选出改进动作'}]}
|
||
|
|
|
||
|
|
|
||
|
|
def capture_worker(monkeypatch):
|
||
|
|
finished = threading.Event()
|
||
|
|
failures = []
|
||
|
|
original = tm._background_assessment
|
||
|
|
|
||
|
|
def worker(*args):
|
||
|
|
try:
|
||
|
|
original(*args)
|
||
|
|
except BaseException as exc:
|
||
|
|
failures.append(exc)
|
||
|
|
finally:
|
||
|
|
finished.set()
|
||
|
|
|
||
|
|
monkeypatch.setattr(tm, '_background_assessment', worker)
|
||
|
|
return finished, failures
|
||
|
|
|
||
|
|
|
||
|
|
def test_modes_validate_and_default_to_full():
|
||
|
|
assert tm.AssessmentInput(version=1).mode == 'full'
|
||
|
|
assert tm.AssessmentInput(version=1, mode='quick').mode == 'quick'
|
||
|
|
with pytest.raises(ValidationError):
|
||
|
|
tm.AssessmentInput(version=1, mode='everything')
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize('mode', ['quick', 'full'])
|
||
|
|
def test_background_returns_progress_and_rejects_duplicate_or_edit_without_wait(assessment_env, monkeypatch, mode):
|
||
|
|
env = assessment_env
|
||
|
|
entered, release = threading.Event(), threading.Event()
|
||
|
|
finished, failures = capture_worker(monkeypatch)
|
||
|
|
|
||
|
|
def ai(session, row, schema, instructions, evidence):
|
||
|
|
assert schema is tm.Plan
|
||
|
|
entered.set()
|
||
|
|
assert release.wait(5), 'test did not release model'
|
||
|
|
return neutral_plan(evidence)
|
||
|
|
|
||
|
|
monkeypatch.setattr(tm, 'model_json', ai)
|
||
|
|
try:
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
started = time.monotonic()
|
||
|
|
accepted = tm.assess(env.draft_id, tm.AssessmentInput(version=1, mode=mode),
|
||
|
|
request_for(session, env), session, background=True)
|
||
|
|
assert time.monotonic() - started < 1
|
||
|
|
assert accepted['status'] == 'assessing' and accepted['coverage']['mode'] == mode
|
||
|
|
assert entered.wait(5)
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
request = request_for(session, env)
|
||
|
|
row = tm.get_draft(env.draft_id, request, session)
|
||
|
|
assert row.status == 'assessing' and row.coverage['stage'] == 'planning'
|
||
|
|
assert row.evidence and row.coverage['evidence_count'] >= 1
|
||
|
|
if mode == 'full':
|
||
|
|
assert row.coverage['scanned_chunks'] >= 1
|
||
|
|
started = time.monotonic()
|
||
|
|
with pytest.raises(HTTPException) as duplicate:
|
||
|
|
tm.assess(row.id, tm.AssessmentInput(version=row.version, mode=mode), request, session, True)
|
||
|
|
assert duplicate.value.status_code == 409
|
||
|
|
with pytest.raises(HTTPException) as editing:
|
||
|
|
tm.update_draft(row.id, tm.UpdateInput(version=row.version, topic='调整时间安排', kb_ids=[env.kb_id]), request, session)
|
||
|
|
assert editing.value.status_code == 409
|
||
|
|
assert time.monotonic() - started < 1
|
||
|
|
finally:
|
||
|
|
release.set()
|
||
|
|
assert finished.wait(5), 'assessment worker failed to finish'
|
||
|
|
assert not failures
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
row = session.get(tm.TrainingMap, env.draft_id)
|
||
|
|
assert row.status == 'proposed'
|
||
|
|
assert row.coverage['stage'] == 'complete' and row.coverage['read_chunks'] >= 1
|
||
|
|
assert not tm._assessment_jobs
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize('revoked', ['user', 'kb'])
|
||
|
|
def test_background_rechecks_permissions_before_publishing(assessment_env, monkeypatch, revoked):
|
||
|
|
env = assessment_env
|
||
|
|
entered, release = threading.Event(), threading.Event()
|
||
|
|
finished, failures = capture_worker(monkeypatch)
|
||
|
|
|
||
|
|
def ai(session, row, schema, instructions, evidence):
|
||
|
|
entered.set()
|
||
|
|
assert release.wait(5)
|
||
|
|
return neutral_plan(evidence)
|
||
|
|
|
||
|
|
monkeypatch.setattr(tm, 'model_json', ai)
|
||
|
|
try:
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
tm.assess(env.draft_id, tm.AssessmentInput(version=1, mode='quick'), request_for(session, env), session, True)
|
||
|
|
assert entered.wait(5)
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
target = session.get(AppUser, env.user_id) if revoked == 'user' else session.get(KnowledgeBaseGrant, env.grant_id)
|
||
|
|
target.enabled = False
|
||
|
|
session.add(target); session.commit()
|
||
|
|
finally:
|
||
|
|
release.set()
|
||
|
|
assert finished.wait(5)
|
||
|
|
assert not failures
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
row = session.get(tm.TrainingMap, env.draft_id)
|
||
|
|
assert row.status == 'failed' and row.plan == {}
|
||
|
|
assert row.coverage['stage'] == 'failed'
|
||
|
|
assert row.coverage['evidence_count'] == 1
|
||
|
|
|
||
|
|
|
||
|
|
def test_interrupted_background_is_retryable_and_keeps_progress(assessment_env, monkeypatch):
|
||
|
|
env = assessment_env
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
row = session.get(tm.TrainingMap, env.draft_id)
|
||
|
|
row.status = 'assessing'
|
||
|
|
row.coverage = {'mode': 'full', 'started_epoch': time.time() - 60,
|
||
|
|
'stage': 'reading', 'read_chunks': 25, 'scanned_chunks': 30}
|
||
|
|
session.add(row); session.commit()
|
||
|
|
row = tm.get_draft(row.id, request_for(session, env), session)
|
||
|
|
assert row.status == 'failed' and row.coverage['read_chunks'] == 25
|
||
|
|
monkeypatch.setattr(tm, 'model_json', lambda session, row, schema, instructions, evidence: neutral_plan(evidence))
|
||
|
|
restarted = tm.assess(row.id, tm.AssessmentInput(version=row.version, mode='quick'), request_for(session, env), session)
|
||
|
|
assert restarted.status == 'proposed'
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize('mode', ['quick', 'full'])
|
||
|
|
def test_model_failure_retains_persisted_coverage(assessment_env, monkeypatch, mode):
|
||
|
|
env = assessment_env
|
||
|
|
def ai(*args, **kwargs):
|
||
|
|
raise ValueError('neutral model failure')
|
||
|
|
monkeypatch.setattr(tm, 'model_json', ai)
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
with pytest.raises(HTTPException):
|
||
|
|
tm.assess(env.draft_id, tm.AssessmentInput(version=1, mode=mode), request_for(session, env), session)
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
row = session.get(tm.TrainingMap, env.draft_id)
|
||
|
|
assert row.status == 'failed' and row.coverage['stage'] == 'failed'
|
||
|
|
assert row.coverage['mode'] == mode
|
||
|
|
assert row.coverage['evidence_count'] == 1
|
||
|
|
assert row.evidence
|
||
|
|
|
||
|
|
|
||
|
|
def test_cancel_during_final_model_prevents_publication(assessment_env, monkeypatch):
|
||
|
|
env = assessment_env
|
||
|
|
entered, release = threading.Event(), threading.Event()
|
||
|
|
finished, failures = capture_worker(monkeypatch)
|
||
|
|
def ai(session, row, schema, instructions, evidence):
|
||
|
|
entered.set()
|
||
|
|
assert release.wait(5)
|
||
|
|
return neutral_plan(evidence)
|
||
|
|
monkeypatch.setattr(tm, 'model_json', ai)
|
||
|
|
try:
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
accepted = tm.assess(env.draft_id, tm.AssessmentInput(version=1, mode='quick'), request_for(session, env), session, True)
|
||
|
|
assert entered.wait(5)
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
started = time.monotonic()
|
||
|
|
cancelled = tm.cancel_assessment(env.draft_id, tm.VersionInput(version=accepted['version']), request_for(session, env), session)
|
||
|
|
assert cancelled.assessment_cancel_requested
|
||
|
|
assert time.monotonic() - started < 1
|
||
|
|
finally:
|
||
|
|
release.set()
|
||
|
|
assert finished.wait(5)
|
||
|
|
assert not failures
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
row = session.get(tm.TrainingMap, env.draft_id)
|
||
|
|
assert row.status == 'failed' and row.plan == {}
|
||
|
|
assert row.coverage['stage'] == 'cancelled'
|
||
|
|
|
||
|
|
|
||
|
|
def test_sources_require_current_version_owner_and_live_kb_grant(assessment_env, monkeypatch):
|
||
|
|
from app import training_corpus
|
||
|
|
env = assessment_env
|
||
|
|
monkeypatch.setattr(tm, 'model_json', lambda session, row, schema, instructions, evidence: neutral_plan(evidence))
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
request = request_for(session, env)
|
||
|
|
row = tm.assess(env.draft_id, tm.AssessmentInput(version=1, mode='full'), request, session)
|
||
|
|
full_version, evidence_id = row.version, row.evidence[0]['id']
|
||
|
|
source = tm.evidence_sources(row.id, evidence_id, request, version=full_version, offset=0, limit=20, s=session)
|
||
|
|
assert source['total'] == 1 and source['items'][0]['kb_id'] == env.kb_id
|
||
|
|
assert '连续三天记录' in source['items'][0]['text']
|
||
|
|
stranger = SimpleNamespace(state=SimpleNamespace(user=AppUser(id=90, username='stranger', role='admin')))
|
||
|
|
with pytest.raises(HTTPException) as denied:
|
||
|
|
tm.evidence_sources(row.id, evidence_id, stranger, version=full_version, offset=0, limit=20, s=session)
|
||
|
|
assert denied.value.status_code == 404
|
||
|
|
# Switching to quick mode creates a new version, without overwriting the full snapshot.
|
||
|
|
row = tm.assess(row.id, tm.AssessmentInput(version=full_version, mode='quick'), request, session)
|
||
|
|
assert (training_corpus.CORPUS_DIR / str(row.id) / f'v{full_version}.sqlite').is_file()
|
||
|
|
with pytest.raises(HTTPException) as stale:
|
||
|
|
tm.evidence_sources(row.id, evidence_id, request, version=full_version, offset=0, limit=20, s=session)
|
||
|
|
assert stale.value.status_code == 409
|
||
|
|
current_version = row.version
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
grant = session.get(KnowledgeBaseGrant, env.grant_id)
|
||
|
|
grant.enabled = False
|
||
|
|
session.add(grant); session.commit()
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
with pytest.raises(HTTPException) as revoked:
|
||
|
|
tm.evidence_sources(env.draft_id, 'E1', request_for(session, env), version=current_version, offset=0, limit=20, s=session)
|
||
|
|
assert revoked.value.status_code == 404
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize('stop', ['user', 'kb', 'cancel'])
|
||
|
|
def test_full_reader_stops_before_another_batch_after_revocation_or_cancel(assessment_env, monkeypatch, stop):
|
||
|
|
env = assessment_env
|
||
|
|
# Three long, distinct entries force multiple real corpus batches.
|
||
|
|
store._write_fallback(env.kb_id, [
|
||
|
|
{'id': f'long-{i}', 'text': f'第{i}份时间记录。' + (f'记录第{i}类任务的时长和来源。' * 700),
|
||
|
|
'meta': {'file': 'lesson.txt'}, 'embedding': [1.0, 0.0]} for i in range(1, 4)])
|
||
|
|
entered, release = threading.Event(), threading.Event()
|
||
|
|
finished, failures = capture_worker(monkeypatch)
|
||
|
|
calls = []
|
||
|
|
|
||
|
|
def ai(session, row, schema, instructions, evidence):
|
||
|
|
calls.append(schema.__name__)
|
||
|
|
assert schema is not tm.Plan, 'should stop before planning'
|
||
|
|
entered.set()
|
||
|
|
assert release.wait(5)
|
||
|
|
return {'text': '记录任务时长、任务类别和干扰来源，再用记录对照并选择改进动作。'}
|
||
|
|
|
||
|
|
monkeypatch.setattr(tm, 'model_json', ai)
|
||
|
|
try:
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
accepted = tm.assess(env.draft_id, tm.AssessmentInput(version=1), request_for(session, env), session, True)
|
||
|
|
assert entered.wait(5)
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
if stop == 'cancel':
|
||
|
|
tm.cancel_assessment(env.draft_id, tm.VersionInput(version=accepted['version']), request_for(session, env), session)
|
||
|
|
else:
|
||
|
|
target = session.get(AppUser, env.user_id) if stop == 'user' else session.get(KnowledgeBaseGrant, env.grant_id)
|
||
|
|
target.enabled = False
|
||
|
|
session.add(target); session.commit()
|
||
|
|
finally:
|
||
|
|
release.set()
|
||
|
|
assert finished.wait(5)
|
||
|
|
assert not failures
|
||
|
|
assert len(calls) == 1
|
||
|
|
with Session(env.engine) as session:
|
||
|
|
row = session.get(tm.TrainingMap, env.draft_id)
|
||
|
|
assert row.status == 'failed' and row.plan == {}
|
||
|
|
assert row.coverage['scanned_chunks'] == 3
|
||
|
|
assert row.coverage['stage'] == ('cancelled' if stop == 'cancel' else 'failed')
|