generated from coulomb/repo-seed
Implemented durable workflow/job foundation
This commit is contained in:
345
tests/test_workflow_service.py
Normal file
345
tests/test_workflow_service.py
Normal file
@@ -0,0 +1,345 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from kontextual_engine import (
|
||||
Actor,
|
||||
ActorType,
|
||||
AssetRegistryService,
|
||||
AssetRepresentation,
|
||||
Classification,
|
||||
InMemoryAssetRegistryRepository,
|
||||
OperationContext,
|
||||
RepresentationKind,
|
||||
Sensitivity,
|
||||
SQLiteAssetRegistryRepository,
|
||||
TransformationRunStatus,
|
||||
TransformationService,
|
||||
ValidationError,
|
||||
WorkflowInputDefinition,
|
||||
WorkflowInputKind,
|
||||
WorkflowInvocation,
|
||||
WorkflowRunStatus,
|
||||
WorkflowService,
|
||||
WorkflowStepDefinition,
|
||||
WorkflowStepRunStatus,
|
||||
WorkflowTemplate,
|
||||
)
|
||||
|
||||
|
||||
def test_workflow_template_registration_persists_input_kinds_and_rejects_bad_dependencies() -> None:
|
||||
repo = InMemoryAssetRegistryRepository()
|
||||
service = WorkflowService(repo)
|
||||
context = operation_context()
|
||||
|
||||
template = service.register_template(rich_input_template(), context)
|
||||
|
||||
assert template.created_by == "user-test"
|
||||
assert [item.kind for item in repo.get_workflow_template("workflow-rich-inputs").inputs] == [
|
||||
WorkflowInputKind.ASSET,
|
||||
WorkflowInputKind.COLLECTION,
|
||||
WorkflowInputKind.QUERY,
|
||||
WorkflowInputKind.SOURCE_EVENT,
|
||||
WorkflowInputKind.PAYLOAD,
|
||||
]
|
||||
assert repo.list_audit_events(target="workflow_template:workflow-rich-inputs")[0].operation == (
|
||||
"workflow.template.register"
|
||||
)
|
||||
|
||||
with pytest.raises(ValidationError) as missing_dependency:
|
||||
service.register_template(
|
||||
WorkflowTemplate(
|
||||
template_id="workflow-bad",
|
||||
name="Bad Workflow",
|
||||
steps=(
|
||||
WorkflowStepDefinition(
|
||||
step_id="late",
|
||||
operation_id="structured_view",
|
||||
depends_on=("missing",),
|
||||
),
|
||||
),
|
||||
),
|
||||
context,
|
||||
)
|
||||
assert missing_dependency.value.details["diagnostics"][0]["code"] == "workflow.dependency_missing"
|
||||
|
||||
with pytest.raises(ValidationError) as cycle:
|
||||
service.register_template(
|
||||
WorkflowTemplate(
|
||||
template_id="workflow-cycle",
|
||||
name="Cycle",
|
||||
steps=(
|
||||
WorkflowStepDefinition(
|
||||
step_id="a",
|
||||
operation_id="structured_view",
|
||||
depends_on=("b",),
|
||||
),
|
||||
WorkflowStepDefinition(
|
||||
step_id="b",
|
||||
operation_id="structured_view",
|
||||
depends_on=("a",),
|
||||
),
|
||||
),
|
||||
),
|
||||
context,
|
||||
)
|
||||
assert cycle.value.details["diagnostics"][0]["code"] == "workflow.dependency_cycle"
|
||||
|
||||
|
||||
def test_workflow_invocation_executes_dependent_transformations_in_order() -> None:
|
||||
repo = InMemoryAssetRegistryRepository()
|
||||
registry = AssetRegistryService(repo)
|
||||
context = operation_context()
|
||||
create_source_asset(registry, context, asset_id="asset-source")
|
||||
workflow = WorkflowService(
|
||||
repo,
|
||||
transformation_service=TransformationService(repo, asset_service=registry),
|
||||
)
|
||||
workflow.register_template(dependent_transformation_template(), context)
|
||||
|
||||
result = workflow.invoke_template(
|
||||
WorkflowInvocation(
|
||||
template_id="workflow-dependent",
|
||||
inputs={"source": {"kind": "asset", "asset_id": "asset-source"}},
|
||||
),
|
||||
context,
|
||||
)
|
||||
|
||||
assert result.success is True
|
||||
assert result.run.status == WorkflowRunStatus.COMPLETED
|
||||
step_runs = {step.step_id: step for step in result.run.step_runs}
|
||||
assert step_runs["view"].status == WorkflowStepRunStatus.COMPLETED
|
||||
assert step_runs["view_again"].status == WorkflowStepRunStatus.COMPLETED
|
||||
assert result.run.output_asset_ids == ("asset-view-1", "asset-view-2")
|
||||
first_transformation = repo.get_transformation_run(step_runs["view"].transformation_run_id)
|
||||
second_transformation = repo.get_transformation_run(step_runs["view_again"].transformation_run_id)
|
||||
assert first_transformation.source_asset_ids == ("asset-source",)
|
||||
assert second_transformation.source_asset_ids == ("asset-view-1",)
|
||||
assert second_transformation.status == TransformationRunStatus.COMPLETED
|
||||
assert [
|
||||
event.operation
|
||||
for event in repo.list_audit_events(target=f"workflow_run:{result.run.run_id}")
|
||||
] == [
|
||||
"workflow.run.queued",
|
||||
"workflow.run.started",
|
||||
"workflow.run.completed",
|
||||
]
|
||||
|
||||
|
||||
def test_workflow_queue_cancel_resume_and_retry_do_not_require_storage_edits() -> None:
|
||||
repo = InMemoryAssetRegistryRepository()
|
||||
registry = AssetRegistryService(repo)
|
||||
context = operation_context()
|
||||
create_source_asset(registry, context, asset_id="asset-source")
|
||||
workflow = WorkflowService(
|
||||
repo,
|
||||
transformation_service=TransformationService(repo, asset_service=registry),
|
||||
)
|
||||
workflow.register_template(single_step_template(), context)
|
||||
|
||||
queued = workflow.queue_template(
|
||||
WorkflowInvocation(
|
||||
template_id="workflow-single",
|
||||
inputs={"source": {"kind": "asset", "asset_id": "asset-source"}},
|
||||
),
|
||||
context,
|
||||
)
|
||||
canceled = workflow.cancel_run(queued.run.run_id, context, reason="operator pause")
|
||||
resumed = workflow.resume_run(canceled.run_id, context)
|
||||
completed = workflow.invoke_template(
|
||||
WorkflowInvocation(
|
||||
template_id="workflow-single",
|
||||
inputs={"source": {"kind": "asset", "asset_id": "asset-source"}},
|
||||
),
|
||||
context,
|
||||
)
|
||||
retry = workflow.retry_run(completed.run.run_id, context)
|
||||
|
||||
assert queued.run.status == WorkflowRunStatus.QUEUED
|
||||
assert canceled.status == WorkflowRunStatus.CANCELED
|
||||
assert {step.status for step in canceled.step_runs} == {WorkflowStepRunStatus.CANCELED}
|
||||
assert resumed.success is False
|
||||
assert resumed.diagnostics[0].code == "workflow.run_not_resumable"
|
||||
assert repo.get_workflow_run(completed.run.run_id).status == WorkflowRunStatus.RETRIED
|
||||
assert retry.run.retry_of_run_id == completed.run.run_id
|
||||
assert retry.run.attempt == 2
|
||||
assert retry.run.status == WorkflowRunStatus.COMPLETED
|
||||
|
||||
|
||||
def test_workflow_continue_failure_finishes_partially_completed() -> None:
|
||||
repo = InMemoryAssetRegistryRepository()
|
||||
registry = AssetRegistryService(repo)
|
||||
context = operation_context()
|
||||
create_source_asset(registry, context, asset_id="asset-source")
|
||||
workflow = WorkflowService(
|
||||
repo,
|
||||
transformation_service=TransformationService(repo, asset_service=registry),
|
||||
)
|
||||
workflow.register_template(partial_template(), context)
|
||||
|
||||
result = workflow.invoke_template(
|
||||
WorkflowInvocation(
|
||||
template_id="workflow-partial",
|
||||
inputs={"source": {"kind": "asset", "asset_id": "asset-source"}},
|
||||
),
|
||||
context,
|
||||
)
|
||||
|
||||
step_runs = {step.step_id: step for step in result.run.step_runs}
|
||||
assert result.success is False
|
||||
assert result.run.status == WorkflowRunStatus.PARTIALLY_COMPLETED
|
||||
assert result.run.output_asset_ids == ("asset-view-success",)
|
||||
assert step_runs["markdown"].status == WorkflowStepRunStatus.SKIPPED
|
||||
assert step_runs["view"].status == WorkflowStepRunStatus.COMPLETED
|
||||
assert result.diagnostics[0].code == "transformation.operation_not_executable"
|
||||
assert repo.list_assets(asset_type="derived_artifact")[0].id == "asset-view-success"
|
||||
|
||||
|
||||
def test_sqlite_workflow_templates_and_runs_survive_reinstantiation(tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "registry.sqlite"
|
||||
repo = SQLiteAssetRegistryRepository(db_path)
|
||||
registry = AssetRegistryService(repo)
|
||||
context = operation_context()
|
||||
create_source_asset(registry, context, asset_id="asset-source")
|
||||
workflow = WorkflowService(
|
||||
repo,
|
||||
transformation_service=TransformationService(repo, asset_service=registry),
|
||||
)
|
||||
workflow.register_template(single_step_template(), context)
|
||||
|
||||
result = workflow.invoke_template(
|
||||
WorkflowInvocation(
|
||||
template_id="workflow-single",
|
||||
inputs={"source": {"kind": "asset", "asset_id": "asset-source"}},
|
||||
),
|
||||
context,
|
||||
)
|
||||
|
||||
reloaded = SQLiteAssetRegistryRepository(db_path)
|
||||
run = reloaded.get_workflow_run(result.run.run_id)
|
||||
|
||||
assert reloaded.get_workflow_template("workflow-single").template_id == "workflow-single"
|
||||
assert run.status == WorkflowRunStatus.COMPLETED
|
||||
assert run.step_runs[0].status == WorkflowStepRunStatus.COMPLETED
|
||||
assert reloaded.list_workflow_runs(template_id="workflow-single")[0].run_id == result.run.run_id
|
||||
assert reloaded.list_representations(asset_id="asset-single-output")[0].kind == RepresentationKind.DERIVED
|
||||
|
||||
|
||||
def rich_input_template() -> WorkflowTemplate:
|
||||
return WorkflowTemplate(
|
||||
template_id="workflow-rich-inputs",
|
||||
name="Rich Inputs",
|
||||
inputs=(
|
||||
WorkflowInputDefinition("source", WorkflowInputKind.ASSET),
|
||||
WorkflowInputDefinition("bundle", WorkflowInputKind.COLLECTION, required=False),
|
||||
WorkflowInputDefinition("query", WorkflowInputKind.QUERY, required=False),
|
||||
WorkflowInputDefinition("event", WorkflowInputKind.SOURCE_EVENT, required=False),
|
||||
WorkflowInputDefinition("payload", WorkflowInputKind.PAYLOAD, required=False),
|
||||
),
|
||||
steps=(
|
||||
WorkflowStepDefinition(
|
||||
step_id="view",
|
||||
operation_id="structured_view",
|
||||
inputs={"source_asset_ids": "$inputs.source"},
|
||||
outputs={"asset_id": "asset-rich-view"},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def dependent_transformation_template() -> WorkflowTemplate:
|
||||
return WorkflowTemplate(
|
||||
template_id="workflow-dependent",
|
||||
name="Dependent Transformations",
|
||||
inputs=(WorkflowInputDefinition("source", WorkflowInputKind.ASSET),),
|
||||
steps=(
|
||||
WorkflowStepDefinition(
|
||||
step_id="view",
|
||||
operation_id="structured_view",
|
||||
inputs={"source_asset_ids": "$inputs.source"},
|
||||
outputs={"asset_id": "asset-view-1", "title": "First View"},
|
||||
),
|
||||
WorkflowStepDefinition(
|
||||
step_id="view_again",
|
||||
operation_id="structured_view",
|
||||
depends_on=("view",),
|
||||
inputs={"source_asset_ids": "$steps.view.output_asset_ids"},
|
||||
outputs={"asset_id": "asset-view-2", "title": "Second View"},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def single_step_template() -> WorkflowTemplate:
|
||||
return WorkflowTemplate(
|
||||
template_id="workflow-single",
|
||||
name="Single Step",
|
||||
inputs=(WorkflowInputDefinition("source", WorkflowInputKind.ASSET),),
|
||||
steps=(
|
||||
WorkflowStepDefinition(
|
||||
step_id="view",
|
||||
operation_id="structured_view",
|
||||
inputs={"source_asset_ids": "$inputs.source"},
|
||||
outputs={"asset_id": "asset-single-output", "title": "Single Output"},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def partial_template() -> WorkflowTemplate:
|
||||
return WorkflowTemplate(
|
||||
template_id="workflow-partial",
|
||||
name="Partial Workflow",
|
||||
inputs=(WorkflowInputDefinition("source", WorkflowInputKind.ASSET),),
|
||||
steps=(
|
||||
WorkflowStepDefinition(
|
||||
step_id="markdown",
|
||||
operation_id="markdown_transform",
|
||||
inputs={"source_asset_ids": "$inputs.source"},
|
||||
outputs={"asset_id": "asset-markdown-output"},
|
||||
failure_behavior="continue",
|
||||
),
|
||||
WorkflowStepDefinition(
|
||||
step_id="view",
|
||||
operation_id="structured_view",
|
||||
inputs={"source_asset_ids": "$inputs.source"},
|
||||
outputs={"asset_id": "asset-view-success"},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def create_source_asset(
|
||||
registry: AssetRegistryService,
|
||||
context: OperationContext,
|
||||
*,
|
||||
asset_id: str,
|
||||
) -> None:
|
||||
registry.create_asset(
|
||||
"Source",
|
||||
Classification(
|
||||
asset_type="document",
|
||||
sensitivity=Sensitivity.INTERNAL,
|
||||
owner="Platform Knowledge",
|
||||
),
|
||||
context,
|
||||
asset_id=asset_id,
|
||||
representations=[
|
||||
AssetRepresentation.from_content(
|
||||
asset_id,
|
||||
RepresentationKind.SOURCE,
|
||||
"text/markdown",
|
||||
"# Source\n",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def operation_context() -> OperationContext:
|
||||
actor = Actor.create(
|
||||
ActorType.HUMAN,
|
||||
actor_id="user-test",
|
||||
display_name="Test User",
|
||||
groups=["engineering"],
|
||||
)
|
||||
return OperationContext.create(actor, correlation_id="corr-test")
|
||||
Reference in New Issue
Block a user