generated from coulomb/repo-seed
346 lines
12 KiB
Python
346 lines
12 KiB
Python
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")
|