windmill/backend/parsers/windmill-parser-wac/tests/python_tests.rs

use windmill_parser_wac::dag::DagNodeType;
use windmill_parser_wac::python::parse_python_workflow;

#[test]
fn test_simple_sequential_workflow() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task
async def extract_data(url: str): ...
@task
async def load_data(data: list): ...

@workflow
async def my_etl(url: str):
    raw = await extract_data(url=url)
    await load_data(data=raw)
    return {"status": "done"}
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    assert_eq!(dag.nodes.len(), 3); // 2 steps + 1 return
    assert_eq!(dag.edges.len(), 2); // step0->step1, step1->return

    // Check params (url — no ctx to skip)
    assert_eq!(dag.params.len(), 1);
    assert_eq!(dag.params[0].name, "url");
    assert_eq!(dag.params[0].typ.as_deref(), Some("str"));

    // Check first step
    match &dag.nodes[0].node_type {
        DagNodeType::Step { name, script } => {
            assert_eq!(name, "extract_data");
            assert_eq!(script, "extract_data");
        }
        _ => panic!("expected Step node"),
    }

    // Check second step
    match &dag.nodes[1].node_type {
        DagNodeType::Step { name, script } => {
            assert_eq!(name, "load_data");
            assert_eq!(script, "load_data");
        }
        _ => panic!("expected Step node"),
    }

    // Check return
    assert!(matches!(dag.nodes[2].node_type, DagNodeType::Return));

    // Check source hash is non-empty
    assert!(!dag.source_hash.is_empty());
}

#[test]
fn test_parallel_workflow() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task
async def extract_data(url: str): ...
@task
async def clean_data(data: list): ...
@task
async def compute_stats(data: list): ...
@task
async def load_to_warehouse(rows: list): ...

@workflow
async def my_etl(url: str):
    raw = await extract_data(url=url)
    cleaned, stats = await asyncio.gather(
        clean_data(data=raw),
        compute_stats(data=raw),
    )
    await load_to_warehouse(rows=cleaned)
    return {"status": "done"}
"#;

    let dag = parse_python_workflow(code).expect("should parse");

    // extract, ParallelStart, clean, stats, ParallelEnd, load, return = 7
    assert_eq!(dag.nodes.len(), 7);

    assert!(matches!(dag.nodes[0].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[1].node_type, DagNodeType::ParallelStart));
    assert!(matches!(dag.nodes[2].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[3].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[4].node_type, DagNodeType::ParallelEnd));
    assert!(matches!(dag.nodes[5].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[6].node_type, DagNodeType::Return));
}

#[test]
fn test_conditional_workflow() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task
async def send_alert(msg: str): ...
@task
async def load_data(): ...

@workflow
async def my_etl(count: int):
    if count > 100:
        await send_alert(msg="large")
    await load_data()
    return {"done": True}
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    // Branch, notify step, load step, return = 4
    assert_eq!(dag.nodes.len(), 4);
    assert!(matches!(dag.nodes[0].node_type, DagNodeType::Branch { .. }));
    assert!(matches!(dag.nodes[1].node_type, DagNodeType::Step { .. }));
}

#[test]
fn test_for_loop_workflow() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task
async def process_item(item: str): ...

@workflow
async def my_etl(items: list):
    for item in items:
        await process_item(item=item)
    return {"done": True}
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    // LoopStart, step, LoopEnd, return = 4
    assert_eq!(dag.nodes.len(), 4);
    assert!(matches!(
        dag.nodes[0].node_type,
        DagNodeType::LoopStart { .. }
    ));
    assert!(matches!(dag.nodes[1].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[2].node_type, DagNodeType::LoopEnd));
}

#[test]
fn test_step_in_try_except() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task
async def extract_data(): ...

@task
async def handle_error(): ...

@workflow
async def my_etl():
    try:
        await extract_data()
    except Exception:
        await handle_error()
"#;

    let dag = parse_python_workflow(code).expect("should parse try/except");
    // Branch(try/except), extract_data, handle_error, merge = 4
    assert_eq!(dag.nodes.len(), 4);
    assert!(matches!(dag.nodes[0].node_type, DagNodeType::Branch { .. }));
    assert_eq!(dag.nodes[0].label, "try");
    assert!(matches!(dag.nodes[1].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[2].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[3].node_type, DagNodeType::Merge));
}

#[test]
fn test_step_in_while() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task
async def poll_status(): ...

@workflow
async def my_etl():
    while True:
        await poll_status()
"#;

    let dag = parse_python_workflow(code).expect("should parse while loop");
    // LoopStart, poll_status, LoopEnd = 3
    assert_eq!(dag.nodes.len(), 3);
    assert!(matches!(
        dag.nodes[0].node_type,
        DagNodeType::LoopStart { .. }
    ));
    assert_eq!(dag.nodes[0].label, "while");
    assert!(matches!(dag.nodes[1].node_type, DagNodeType::Step { .. }));
    assert!(matches!(dag.nodes[2].node_type, DagNodeType::LoopEnd));
}

#[test]
fn test_reject_non_async() {
    let code = r#"
from wmill import workflow

@workflow
def my_etl():
    pass
"#;

    let result = parse_python_workflow(code);
    assert!(result.is_err());
    let errors = result.unwrap_err();
    assert!(errors[0].message.contains("async"));
}

#[test]
fn test_reject_missing_await() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task
async def extract_data(): ...

@workflow
async def my_etl():
    extract_data()
"#;

    let result = parse_python_workflow(code);
    assert!(result.is_err());
    let errors = result.unwrap_err();
    assert!(errors[0].message.contains("awaited"));
}

#[test]
fn test_no_workflow_function() {
    let code = r#"
async def my_func():
    pass
"#;

    let result = parse_python_workflow(code);
    assert!(result.is_err());
    let errors = result.unwrap_err();
    assert!(errors[0].message.contains("No @workflow"));
}

#[test]
fn test_task_with_external_path() {
    let code = r#"
import asyncio
from wmill import workflow, task

@task(path="f/external_script")
async def run_external(x: int): ...

@workflow
async def my_wf(x: int):
    result = await run_external(x=x)
    return result
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    assert_eq!(dag.nodes.len(), 2); // 1 step + 1 return (bare `return` is not a step node but walk_return creates one)
    match &dag.nodes[0].node_type {
        DagNodeType::Step { name, script } => {
            assert_eq!(name, "run_external");
            assert_eq!(script, "f/external_script");
        }
        _ => panic!("expected Step node"),
    }
}

#[test]
fn test_task_script_and_task_flow_py() {
    let code = r#"
from wmill import workflow, task, task_script, task_flow

helper = task_script("./helper.py")
pipeline = task_flow("f/etl/pipeline")

@task()
async def process(x: str) -> str:
    return f"processed: {x}"

@workflow
async def main(x: str):
    a = await process(x=x)
    b = await helper(a=a)
    c = await pipeline(b=b)
    return {"a": a, "b": b, "c": c}
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    assert_eq!(dag.nodes.len(), 4); // 3 steps + 1 return

    match &dag.nodes[1].node_type {
        DagNodeType::Step { name, script } => {
            assert_eq!(name, "helper");
            assert_eq!(script, "./helper.py");
        }
        _ => panic!("expected Step node for task_script"),
    }

    match &dag.nodes[2].node_type {
        DagNodeType::Step { name, script } => {
            assert_eq!(name, "pipeline");
            assert_eq!(script, "f/etl/pipeline");
        }
        _ => panic!("expected Step node for task_flow"),
    }
}

#[test]
fn test_full_template_with_sdk_calls_py() {
    let code = r#"
from wmill import workflow, task, task_script, step, sleep, wait_for_approval, get_resume_urls

helper = task_script("./helper.py")

@task()
async def process(x: str) -> str:
    return f"processed: {x}"

@workflow
async def main(x: str):
    a = await process(x=x)
    b = await helper(a=a)
    urls = await step("get_urls", lambda: get_resume_urls())
    await sleep(1)
    approval = await wait_for_approval(timeout=3600)
    return {"processed": a, "helper_result": b, "approval": approval}
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    // process, helper, step("get_urls"), sleep(1), wait_for_approval, return = 6
    assert_eq!(dag.nodes.len(), 6);

    match &dag.nodes[2].node_type {
        DagNodeType::InlineStep { name } => {
            assert_eq!(name, "get_urls");
        }
        _ => panic!("expected InlineStep node, got {:?}", dag.nodes[2].node_type),
    }

    match &dag.nodes[3].node_type {
        DagNodeType::Sleep { seconds } => {
            assert_eq!(seconds, "1");
        }
        _ => panic!("expected Sleep node, got {:?}", dag.nodes[3].node_type),
    }

    assert!(matches!(
        dag.nodes[4].node_type,
        DagNodeType::WaitForApproval
    ));
    assert!(matches!(dag.nodes[5].node_type, DagNodeType::Return));
}

#[test]
fn test_elif_chain_nests_branches() {
    let code = r#"
from wmill import workflow, task

@task
async def a(): ...
@task
async def b(): ...
@task
async def c(): ...

@workflow
async def main(x: int):
    if x > 2:
        await a()
    elif x == 2:
        await b()
    else:
        await c()
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    let branches: Vec<(&str, usize)> = dag
        .nodes
        .iter()
        .filter_map(|n| match &n.node_type {
            DagNodeType::Branch { condition_source } => Some((condition_source.as_str(), n.line)),
            _ => None,
        })
        .collect();
    assert_eq!(branches, vec![("x > 2", 13), ("x == 2", 15)]);

    // The `elif` branch hangs off the `false` edge of the `if` before it
    let (outer, inner) = (&dag.nodes[0].id, &dag.nodes[2].id);
    assert!(dag
        .edges
        .iter()
        .any(|e| &e.from == outer && &e.to == inner && e.label.as_deref() == Some("false")));
    assert_eq!(
        dag.nodes
            .iter()
            .filter(|n| matches!(n.node_type, DagNodeType::Merge))
            .count(),
        2
    );
}

#[test]
fn test_workflow_with_recent_syntax() {
    let code = r#"
from wmill import workflow, task
lazy import polars

type Rows = list[dict]

@task
async def load[T](rows: T): ...

@workflow
async def main(name: str):
    try:
        await load(t"hello {name}")
    except ValueError, TypeError:
        pass
"#;

    let dag = parse_python_workflow(code).expect("should parse");
    assert!(dag
        .nodes
        .iter()
        .any(|n| matches!(&n.node_type, DagNodeType::Step { name, .. } if name == "load")));
}