361 lines
11 KiB
Python
361 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from wf_artifacts import (
|
|
FileWorkflowArtifactStore,
|
|
WorkflowArtifact,
|
|
WorkflowDeployment,
|
|
)
|
|
|
|
|
|
def artifact(version: int) -> WorkflowArtifact:
|
|
return WorkflowArtifact(
|
|
id="summarize_docs",
|
|
version=version,
|
|
title=f"Summarize Docs v{version}",
|
|
description=None,
|
|
input_schema={"type": "object", "properties": {}},
|
|
output_schema={"type": "object", "properties": {}},
|
|
outcomes=("done",),
|
|
plan={"name": "summarize_docs", "nodes": [], "edges": []},
|
|
)
|
|
|
|
|
|
def test_file_store_round_trips_artifact_versions(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
store.save_artifact(artifact(1))
|
|
store.save_artifact(artifact(2))
|
|
|
|
loaded = store.get_artifact("summarize_docs", 2)
|
|
|
|
assert loaded.id == "summarize_docs"
|
|
assert loaded.version == 2
|
|
assert loaded.title == "Summarize Docs v2"
|
|
|
|
|
|
def test_file_store_resolves_latest_artifact_version(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
store.save_artifact(artifact(1))
|
|
store.save_artifact(artifact(3))
|
|
store.save_artifact(artifact(2))
|
|
|
|
latest = store.resolve_latest("summarize_docs")
|
|
|
|
assert latest.id == "summarize_docs"
|
|
assert latest.version == 3
|
|
|
|
|
|
def test_file_store_round_trips_deployment(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
deployment = WorkflowDeployment(
|
|
id="summarize_docs.personal",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
bindings=[
|
|
{"logical_source": "context7", "concrete_source": "context7.personal"}
|
|
],
|
|
)
|
|
|
|
store.save_deployment(deployment)
|
|
loaded = store.get_deployment("summarize_docs.personal")
|
|
|
|
assert loaded.id == "summarize_docs.personal"
|
|
assert loaded.artifact_id == "summarize_docs"
|
|
assert loaded.binding_map()["context7"] == "context7.personal"
|
|
|
|
|
|
def test_concurrent_deployment_saves_advance_revision_without_lost_updates(
|
|
tmp_path, monkeypatch
|
|
) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
store.save_deployment(
|
|
WorkflowDeployment(
|
|
id="concurrent.personal",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
)
|
|
)
|
|
updates = [
|
|
WorkflowDeployment(
|
|
id="concurrent.personal",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=version,
|
|
)
|
|
for version in range(2, 10)
|
|
]
|
|
|
|
start_gate = threading.Barrier(len(updates))
|
|
read_gate = threading.Barrier(len(updates))
|
|
read_state_lock = threading.Lock()
|
|
active_reads = 0
|
|
max_active_reads = 0
|
|
deployment_path = (store.deployments_dir / "concurrent.personal.json").resolve()
|
|
original_read_text = Path.read_text
|
|
|
|
def coordinated_read(path: Path, *args: Any, **kwargs: Any) -> str:
|
|
nonlocal active_reads, max_active_reads
|
|
if path.resolve() != deployment_path:
|
|
return original_read_text(path, *args, **kwargs)
|
|
with read_state_lock:
|
|
active_reads += 1
|
|
max_active_reads = max(max_active_reads, active_reads)
|
|
try:
|
|
# A broken implementation reaches this gate from every worker
|
|
# after reading the same revision. The real store's lock lets
|
|
# only one worker enter the read/modify/write window, so the
|
|
# gate times out once and later workers proceed immediately.
|
|
try:
|
|
read_gate.wait(timeout=0.5)
|
|
except threading.BrokenBarrierError:
|
|
pass
|
|
return original_read_text(path, *args, **kwargs)
|
|
finally:
|
|
with read_state_lock:
|
|
active_reads -= 1
|
|
|
|
monkeypatch.setattr(Path, "read_text", coordinated_read)
|
|
|
|
def save(deployment: WorkflowDeployment) -> None:
|
|
start_gate.wait(timeout=10)
|
|
store.save_deployment(deployment)
|
|
|
|
with ThreadPoolExecutor(max_workers=len(updates)) as executor:
|
|
list(executor.map(save, updates))
|
|
|
|
assert max_active_reads == 1
|
|
assert store.get_deployment("concurrent.personal").revision == 9
|
|
|
|
|
|
def test_deployment_revision_increments_on_save(tmp_path) -> None:
|
|
artifacts = FileWorkflowArtifactStore(tmp_path)
|
|
artifacts.save_deployment(
|
|
WorkflowDeployment(
|
|
id="dep-1",
|
|
artifact_id="wf-1",
|
|
artifact_version=1,
|
|
bindings=[],
|
|
)
|
|
)
|
|
assert artifacts.get_deployment("dep-1").revision == 1
|
|
artifacts.save_deployment(
|
|
WorkflowDeployment(
|
|
id="dep-1",
|
|
artifact_id="wf-1",
|
|
artifact_version=2,
|
|
bindings=[],
|
|
)
|
|
)
|
|
assert artifacts.get_deployment("dep-1").revision == 2
|
|
|
|
|
|
def test_file_store_loads_legacy_artifact_and_rewrites_canonical_shape(
|
|
tmp_path,
|
|
) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
artifact_dir = store.artifacts_dir / "legacy_capabilities"
|
|
artifact_dir.mkdir(parents=True)
|
|
artifact_path = artifact_dir / "1.json"
|
|
artifact_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"id": "legacy_capabilities",
|
|
"version": 1,
|
|
"title": "Legacy Capabilities",
|
|
"input_schema": {"type": "object", "properties": {}},
|
|
"output_schema": {"type": "object", "properties": {}},
|
|
"outcomes": ["done"],
|
|
"plan": {"name": "legacy_capabilities", "nodes": [], "edges": []},
|
|
"required_capabilities": {
|
|
"demo.echo": {
|
|
"kind": "tool",
|
|
"input_schema_hash": "sha256:input",
|
|
}
|
|
},
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
loaded = store.get_artifact("legacy_capabilities", 1)
|
|
store.save_artifact(loaded)
|
|
rewritten = json.loads(artifact_path.read_text(encoding="utf-8"))
|
|
|
|
required = rewritten["required_capabilities"][0]
|
|
assert required["ref"] == {"source": "demo", "capability_key": "echo"}
|
|
assert required["kind"] == "tool"
|
|
assert "logical_source" not in required
|
|
assert "capability_name" not in required
|
|
|
|
|
|
def test_file_store_loads_legacy_deployment_and_rewrites_canonical_shape(
|
|
tmp_path,
|
|
) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
deployment_path = store.deployments_dir / "legacy_bindings.personal.json"
|
|
deployment_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"id": "legacy_bindings.personal",
|
|
"artifact_id": "legacy_bindings",
|
|
"artifact_version": 1,
|
|
"bindings": {"demo": "demo.personal"},
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
loaded = store.get_deployment("legacy_bindings.personal")
|
|
store.save_deployment(loaded)
|
|
rewritten = json.loads(deployment_path.read_text(encoding="utf-8"))
|
|
|
|
binding = rewritten["bindings"][0]
|
|
assert binding["logical_source"] == "demo"
|
|
assert binding["concrete_source"] == "demo.personal"
|
|
assert loaded.binding_map()["demo"] == "demo.personal"
|
|
|
|
|
|
def test_file_store_lists_deployments_in_id_order(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
store.save_deployment(
|
|
WorkflowDeployment(
|
|
id="summarize_docs.work",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
bindings=[
|
|
{"logical_source": "context7", "concrete_source": "context7.work"}
|
|
],
|
|
)
|
|
)
|
|
store.save_deployment(
|
|
WorkflowDeployment(
|
|
id="summarize_docs.personal",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
bindings=[
|
|
{"logical_source": "context7", "concrete_source": "context7.personal"}
|
|
],
|
|
)
|
|
)
|
|
|
|
deployments = store.list_deployments()
|
|
|
|
assert [deployment.id for deployment in deployments] == [
|
|
"summarize_docs.personal",
|
|
"summarize_docs.work",
|
|
]
|
|
assert deployments[0].binding_map()["context7"] == "context7.personal"
|
|
|
|
|
|
def test_file_store_deletes_deployment(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
store.save_deployment(
|
|
WorkflowDeployment(
|
|
id="summarize_docs.personal",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
bindings=[
|
|
{"logical_source": "context7", "concrete_source": "context7.personal"}
|
|
],
|
|
)
|
|
)
|
|
|
|
store.delete_deployment("summarize_docs.personal")
|
|
|
|
assert store.list_deployments() == []
|
|
|
|
|
|
def test_file_store_rejects_artifact_id_path_traversal(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path / "store")
|
|
bad_artifact = artifact(1).model_copy(update={"id": "../outside"})
|
|
|
|
with pytest.raises(ValueError, match="artifact_id"):
|
|
store.save_artifact(bad_artifact)
|
|
|
|
assert not (tmp_path / "outside").exists()
|
|
|
|
|
|
def test_file_store_rejects_artifact_lookup_path_traversal(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path / "store")
|
|
|
|
with pytest.raises(ValueError, match="artifact_id"):
|
|
store.get_artifact("../outside", 1)
|
|
|
|
|
|
def test_file_store_rejects_deployment_id_path_traversal(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path / "store")
|
|
bad_deployment = WorkflowDeployment(
|
|
id="../outside",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="deployment_id"):
|
|
store.save_deployment(bad_deployment)
|
|
|
|
assert not (tmp_path / "outside.json").exists()
|
|
|
|
|
|
def test_file_store_rejects_deployment_lookup_path_traversal(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path / "store")
|
|
|
|
with pytest.raises(ValueError, match="deployment_id"):
|
|
store.get_deployment("../outside")
|
|
|
|
|
|
def test_file_store_deletes_one_artifact_version(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
store.save_artifact(artifact(1))
|
|
store.save_artifact(artifact(2))
|
|
|
|
store.delete_artifact("summarize_docs", 1)
|
|
|
|
with pytest.raises(KeyError, match="unknown workflow artifact"):
|
|
store.get_artifact("summarize_docs", 1)
|
|
assert store.get_artifact("summarize_docs", 2).version == 2
|
|
|
|
|
|
def test_file_store_delete_artifact_missing_version_raises_key_error(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
|
|
with pytest.raises(KeyError, match="unknown workflow artifact"):
|
|
store.delete_artifact("summarize_docs", 1)
|
|
|
|
|
|
def test_file_store_finds_deployments_for_artifact_version(tmp_path) -> None:
|
|
store = FileWorkflowArtifactStore(tmp_path)
|
|
store.save_deployment(
|
|
WorkflowDeployment(
|
|
id="summarize_docs.work",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
)
|
|
)
|
|
store.save_deployment(
|
|
WorkflowDeployment(
|
|
id="summarize_docs.personal",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=1,
|
|
)
|
|
)
|
|
store.save_deployment(
|
|
WorkflowDeployment(
|
|
id="summarize_docs.v2",
|
|
artifact_id="summarize_docs",
|
|
artifact_version=2,
|
|
)
|
|
)
|
|
|
|
blockers = store.deployments_for_artifact("summarize_docs", 1)
|
|
|
|
assert [deployment.id for deployment in blockers] == [
|
|
"summarize_docs.personal",
|
|
"summarize_docs.work",
|
|
]
|