90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
from types import SimpleNamespace
|
|
|
|
from deerflow.agents.middlewares.artifact_reconcile_middleware import (
|
|
ArtifactReconcileMiddleware,
|
|
)
|
|
from deerflow.agents.thread_state import ARTIFACTS_REPLACE_SENTINEL
|
|
|
|
|
|
def test_before_model_prunes_missing_outputs_artifacts(tmp_path):
|
|
outputs_dir = tmp_path / "outputs"
|
|
outputs_dir.mkdir()
|
|
existing = outputs_dir / "keep.md"
|
|
existing.write_text("ok", encoding="utf-8")
|
|
|
|
middleware = ArtifactReconcileMiddleware()
|
|
state = {
|
|
"thread_data": {"outputs_path": str(outputs_dir)},
|
|
"artifacts": [
|
|
"/mnt/user-data/outputs/keep.md",
|
|
"/mnt/user-data/outputs/missing.md",
|
|
],
|
|
}
|
|
|
|
result = middleware.before_model(state, runtime=SimpleNamespace(context={}))
|
|
|
|
assert result == {
|
|
"artifacts": [ARTIFACTS_REPLACE_SENTINEL, "/mnt/user-data/outputs/keep.md"]
|
|
}
|
|
|
|
|
|
def test_before_model_returns_none_when_no_changes(tmp_path):
|
|
outputs_dir = tmp_path / "outputs"
|
|
outputs_dir.mkdir()
|
|
existing = outputs_dir / "keep.md"
|
|
existing.write_text("ok", encoding="utf-8")
|
|
|
|
middleware = ArtifactReconcileMiddleware()
|
|
state = {
|
|
"thread_data": {"outputs_path": str(outputs_dir)},
|
|
"artifacts": ["/mnt/user-data/outputs/keep.md"],
|
|
}
|
|
|
|
result = middleware.before_model(state, runtime=SimpleNamespace(context={}))
|
|
|
|
assert result is None
|
|
|
|
|
|
def test_before_model_adds_unpresented_outputs_files(tmp_path):
|
|
outputs_dir = tmp_path / "outputs"
|
|
outputs_dir.mkdir()
|
|
existing = outputs_dir / "keep.md"
|
|
existing.write_text("ok", encoding="utf-8")
|
|
extra = outputs_dir / "extra.md"
|
|
extra.write_text("ok", encoding="utf-8")
|
|
|
|
middleware = ArtifactReconcileMiddleware()
|
|
state = {
|
|
"thread_data": {"outputs_path": str(outputs_dir)},
|
|
"artifacts": ["/mnt/user-data/outputs/keep.md"],
|
|
}
|
|
|
|
result = middleware.before_model(state, runtime=SimpleNamespace(context={}))
|
|
|
|
assert result == {
|
|
"artifacts": [
|
|
ARTIFACTS_REPLACE_SENTINEL,
|
|
"/mnt/user-data/outputs/keep.md",
|
|
"/mnt/user-data/outputs/extra.md",
|
|
]
|
|
}
|
|
|
|
|
|
def test_before_model_discovers_outputs_when_artifacts_empty(tmp_path):
|
|
outputs_dir = tmp_path / "outputs"
|
|
outputs_dir.mkdir()
|
|
report = outputs_dir / "report.md"
|
|
report.write_text("ok", encoding="utf-8")
|
|
|
|
middleware = ArtifactReconcileMiddleware()
|
|
state = {
|
|
"thread_data": {"outputs_path": str(outputs_dir)},
|
|
"artifacts": [],
|
|
}
|
|
|
|
result = middleware.before_model(state, runtime=SimpleNamespace(context={}))
|
|
|
|
assert result == {
|
|
"artifacts": [ARTIFACTS_REPLACE_SENTINEL, "/mnt/user-data/outputs/report.md"]
|
|
}
|