312 lines
11 KiB
Python
312 lines
11 KiB
Python
import pytest
|
|
|
|
from tht.corpus.chunk import ChunkPolicy
|
|
from tht.corpus.pipeline import CorpusPipeline, PipelineError
|
|
from tht.corpus.store import CorpusStore
|
|
from tht.ports.evidence import AcquiredDocument, SourceObject
|
|
from tht.ports.vector import VectorCapabilities
|
|
|
|
|
|
class Source:
|
|
def __init__(self, documents):
|
|
self.documents = documents
|
|
self.acquire_calls = []
|
|
|
|
def discover(self):
|
|
return [item[0] for item in self.documents]
|
|
|
|
def acquire(self, item):
|
|
self.acquire_calls.append(item.source_id)
|
|
payload = next(payload for source, payload in self.documents if source.source_id == item.source_id)
|
|
if isinstance(payload, Exception):
|
|
raise payload
|
|
return AcquiredDocument(source=item, content=payload.encode())
|
|
|
|
|
|
class Embedder:
|
|
def __init__(self, dim=3, fail=False):
|
|
self.dim = dim
|
|
self.fail = fail
|
|
self.calls = []
|
|
|
|
def embed_documents(self, texts):
|
|
self.calls.extend(texts)
|
|
if self.fail:
|
|
raise RuntimeError("embed failed")
|
|
return [[float(i) for i in range(self.dim)] for _ in texts]
|
|
|
|
|
|
class Vectors:
|
|
capabilities = VectorCapabilities(search=True, existing_hashes=True, upsert=True)
|
|
|
|
def __init__(self, fail=False):
|
|
self.fail = fail
|
|
self.records = []
|
|
|
|
def upsert(self, collection, records):
|
|
self.records.extend(records[:1] if self.fail else records)
|
|
if self.fail:
|
|
raise RuntimeError("partial write")
|
|
return len(records)
|
|
|
|
def existing_hashes(self, collection, kinds):
|
|
return {
|
|
value.record.id: value.content_hash for value in self.records
|
|
}
|
|
|
|
def delete_generation(self, collection, generation):
|
|
self.records = [
|
|
value for value in self.records
|
|
if value.record.metadata["vector_generation"] != generation
|
|
]
|
|
return 0
|
|
|
|
|
|
def item(name, fingerprint):
|
|
return SourceObject(
|
|
source_id=f"fs:{name}", uri=f"file:///safe/{name}.md", fingerprint=f"sha256:{fingerprint}"
|
|
)
|
|
|
|
|
|
def pipeline(tmp_path, source, *, embedder=None, vectors=None, model="model-a", policy=None):
|
|
return CorpusPipeline(
|
|
store=CorpusStore(tmp_path / "corpus"), sources=[source],
|
|
embedder=embedder or Embedder(), vector_store=vectors or Vectors(),
|
|
embedding_model=model, embedding_dimensions=3,
|
|
chunk_policy=policy or ChunkPolicy(version="chunk-v1", max_chars=100),
|
|
pipeline_version="evidence-v1",
|
|
)
|
|
|
|
|
|
def test_unchanged_documents_skip_acquire_normalize_chunk_and_embed(tmp_path):
|
|
one = item("one", "a")
|
|
first_source = Source([(one, "hello")])
|
|
first = pipeline(tmp_path, first_source)
|
|
first.run()
|
|
second_source = Source([(one, "ignored")])
|
|
second_embedder = Embedder()
|
|
result = pipeline(tmp_path, second_source, embedder=second_embedder).run()
|
|
assert result.unchanged == ("fs:one",)
|
|
assert second_source.acquire_calls == []
|
|
assert second_embedder.calls == []
|
|
|
|
|
|
def test_removed_documents_are_marked_and_absent_from_new_manifest(tmp_path):
|
|
one, two = item("one", "a"), item("two", "b")
|
|
pipeline(tmp_path, Source([(one, "one"), (two, "two")])).run()
|
|
result = pipeline(tmp_path, Source([(one, "one")])).run()
|
|
assert result.removed == ("fs:two",)
|
|
assert {doc.source_id for doc in result.manifest.documents} == {"fs:one"}
|
|
|
|
|
|
def test_model_or_chunk_policy_change_forces_full_rebuild(tmp_path):
|
|
one = item("one", "a")
|
|
pipeline(tmp_path, Source([(one, "hello")])).run()
|
|
source = Source([(one, "hello")])
|
|
changed = pipeline(tmp_path, source, model="model-b").run()
|
|
assert changed.changed == ("fs:one",)
|
|
assert source.acquire_calls == ["fs:one"]
|
|
|
|
|
|
def test_partial_vector_failure_never_changes_active_or_exposes_generation(tmp_path):
|
|
one = item("one", "a")
|
|
good = pipeline(tmp_path, Source([(one, "old")]))
|
|
old = good.run().generation
|
|
changed = item("one", "b")
|
|
vectors = Vectors(fail=True)
|
|
broken = pipeline(tmp_path, Source([(changed, "new")]), vectors=vectors)
|
|
with pytest.raises(PipelineError):
|
|
broken.run()
|
|
assert broken.store.active_generation() == old
|
|
assert vectors.records[0].record.metadata["vector_generation"] != old
|
|
|
|
|
|
def test_dimension_mismatch_fails_before_vector_write_and_publish(tmp_path):
|
|
one = item("one", "a")
|
|
vectors = Vectors()
|
|
candidate = pipeline(tmp_path, Source([(one, "hello")]), embedder=Embedder(dim=2), vectors=vectors)
|
|
with pytest.raises(PipelineError, match="dimension"):
|
|
candidate.run()
|
|
assert vectors.records == []
|
|
assert candidate.store.active_generation() is None
|
|
|
|
|
|
def test_dry_run_and_failed_acquire_never_change_active(tmp_path):
|
|
one = item("one", "a")
|
|
active = pipeline(tmp_path, Source([(one, "old")])).run().generation
|
|
changed = item("one", "b")
|
|
dry = pipeline(tmp_path, Source([(changed, "new")])).run(dry_run=True)
|
|
assert dry.published is False
|
|
assert dry.generation is None
|
|
assert dry.manifest.documents[0].content == "old"
|
|
with pytest.raises(PipelineError):
|
|
pipeline(tmp_path, Source([(changed, RuntimeError("boom"))])).run()
|
|
assert CorpusStore(tmp_path / "corpus").active_generation() == active
|
|
|
|
|
|
def test_job_pipeline_uses_ordered_plan_and_returns_run_id(tmp_path):
|
|
one = item("one", "a")
|
|
candidate = pipeline(tmp_path, Source([(one, "hello")]))
|
|
result = candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
)
|
|
assert result.status == "succeeded"
|
|
assert result.run_id and len(result.run_id) == 32
|
|
checkpoint = tmp_path / ".tht-jobs" / "evidence" / "runs" / result.run_id / "checkpoint.json"
|
|
payload = __import__("json").loads(checkpoint.read_text())
|
|
assert [stage["name"] for stage in payload["stages"]] == [
|
|
"discover", "acquire_normalize_chunk", "embed", "vector_upsert",
|
|
"stage_validate", "publish", "retention_cleanup",
|
|
]
|
|
|
|
|
|
def test_job_pipeline_dry_run_only_discovers_and_reports_changes(tmp_path):
|
|
one = item("one", "a")
|
|
source = Source([(one, "hello")])
|
|
embedder = Embedder()
|
|
vectors = Vectors()
|
|
result = pipeline(tmp_path, source, embedder=embedder, vectors=vectors).run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
dry_run=True,
|
|
)
|
|
assert result.changed == ("fs:one",)
|
|
assert source.acquire_calls == []
|
|
assert embedder.calls == []
|
|
assert vectors.records == []
|
|
assert result.generation is None and result.published is False
|
|
|
|
|
|
@pytest.mark.parametrize("crash_stage", [
|
|
"discover", "acquire_normalize_chunk", "embed", "vector_upsert",
|
|
"stage_validate", "publish", "retention_cleanup",
|
|
])
|
|
def test_job_pipeline_crash_after_each_stage_resumes_without_duplicate_effects(tmp_path, crash_stage):
|
|
one = item("one", "a")
|
|
source = Source([(one, "hello")])
|
|
embedder = Embedder()
|
|
vectors = Vectors()
|
|
candidate = pipeline(tmp_path, source, embedder=embedder, vectors=vectors)
|
|
|
|
class Crash(BaseException):
|
|
pass
|
|
|
|
def fault(_context, stage):
|
|
if stage == crash_stage:
|
|
raise Crash()
|
|
|
|
with pytest.raises(Crash):
|
|
candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
after_stage_return=fault,
|
|
)
|
|
runs = tmp_path / ".tht-jobs" / "evidence" / "runs"
|
|
crashed = next(runs.iterdir()).name
|
|
result = candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
resume_run_id=crashed,
|
|
)
|
|
assert result.status == "succeeded"
|
|
assert source.acquire_calls == ["fs:one"]
|
|
assert len(embedder.calls) == 1
|
|
assert len(vectors.records) == 1
|
|
|
|
|
|
def test_job_pipeline_raw_upsert_failure_compensates_and_resumes_with_new_generation(tmp_path):
|
|
one = item("one", "a")
|
|
vectors = Vectors(fail=True)
|
|
candidate = pipeline(tmp_path, Source([(one, "hello")]), vectors=vectors)
|
|
first = candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
)
|
|
assert first.status == "failed"
|
|
assert vectors.records == []
|
|
old_generation = first.generation
|
|
vectors.fail = False
|
|
resumed = candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
resume_run_id=first.run_id,
|
|
)
|
|
assert resumed.status == "succeeded", resumed
|
|
assert resumed.generation != old_generation
|
|
assert candidate.store.active_generation() == resumed.generation
|
|
|
|
|
|
def test_job_pipeline_raw_stage_failure_compensates_vectors_and_resumes(tmp_path, monkeypatch):
|
|
one = item("one", "a")
|
|
vectors = Vectors()
|
|
candidate = pipeline(tmp_path, Source([(one, "hello")]), vectors=vectors)
|
|
real_stage = candidate.store.stage
|
|
calls = 0
|
|
|
|
def fail_once(*args, **kwargs):
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise OSError("raw stage failure")
|
|
return real_stage(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(candidate.store, "stage", fail_once)
|
|
first = candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
)
|
|
assert first.status == "failed" and vectors.records == []
|
|
resumed = candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
resume_run_id=first.run_id,
|
|
)
|
|
assert resumed.status == "succeeded", resumed
|
|
|
|
|
|
@pytest.mark.parametrize("stage,filename", [
|
|
("discover", "plan.json"),
|
|
("acquire_normalize_chunk", "manifest.json"),
|
|
("embed", "embeddings.json"),
|
|
])
|
|
@pytest.mark.parametrize("mutation", ["missing", "tampered"])
|
|
def test_job_pipeline_rejects_corrupt_required_artifacts_before_resume(
|
|
tmp_path, stage, filename, mutation,
|
|
):
|
|
one = item("one", "a")
|
|
candidate = pipeline(tmp_path, Source([(one, "hello")]))
|
|
|
|
class Crash(BaseException):
|
|
pass
|
|
|
|
with pytest.raises(Crash):
|
|
candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
after_stage_return=lambda _context, name: (
|
|
(_ for _ in ()).throw(Crash()) if name == stage else None
|
|
),
|
|
)
|
|
runs = tmp_path / ".tht-jobs" / "evidence" / "runs"
|
|
crashed = next(runs.iterdir())
|
|
target = crashed / "artifacts" / filename
|
|
target.unlink() if mutation == "missing" else target.write_text("tampered")
|
|
from tht.jobs.runner import CorruptCheckpointError
|
|
with pytest.raises(CorruptCheckpointError, match="artifact"):
|
|
candidate.run_as_job(
|
|
workspace_id="demo", workspace_root=tmp_path,
|
|
config_fingerprint="sha256:" + "1" * 64,
|
|
input_fingerprint="sha256:" + "2" * 64,
|
|
resume_run_id=crashed.name,
|
|
)
|