tai/taicpp/tests/mplex_flow_gen.py

"""Corpus generator for mplex_flow_test.cpp: runs the real tai.py validate_flow_spec and
PlexerFlow over generated specs and writes the results as JSON lines.

usage: python3 -I mplex_flow_gen.py <repo root> <out file>
"""
import copy
import importlib.util
import json
import random
import sys

root, out_path = sys.argv[1], sys.argv[2]
spec = importlib.util.spec_from_file_location("tai_mod", root + "/tai.py")
tai = importlib.util.module_from_spec(spec)
sys.modules["tai_mod"] = tai
spec.loader.exec_module(tai)

rng = random.Random(20260101)

BASES = [
    {"v": 1, "title": "Deploy", "timeout": 300, "start": "main", "screens": [
        {"id": "main", "title": "Where", "body": "Pick a target.",
         "fields": [
             {"id": "env", "type": "choice", "label": "Environment", "options": ["staging", "prod"], "default": 0},
             {"id": "tag", "type": "text", "label": "Tag", "initial": "v1", "show_if": {"env": "prod"}},
             {"id": "note", "type": "info", "text": "Prod deploys page the team."}],
         "actions": [{"id": "go", "label": "Deploy"}, {"id": "back", "label": "Back"}],
         "next": {"go": "$done", "back": "$cancel"}}]},
    {"screens": [{"id": "a", "body": "hello"}, {"id": "b", "title": "B", "fields": [
        {"id": "x", "type": "choice", "options": [{"label": "One", "value": 1}, {"label": "Two", "value": 2.5}, "three"]},
        {"id": "y", "type": "text", "label": "Name é中", "show_if": {"x": [2.5, "three"]}}],
        "next": "c"}, {"id": "c", "title": "C", "show_if": {"x": 1}, "body": "only one"},
        {"id": "d", "actions": [{"id": "ok", "label": "OK"}, {"id": "again", "label": "Again"}],
         "next": {"again": "a"}}]},
    {"title": "T", "timeout": 5, "screens": [{"id": "only", "fields": [{"id": "i", "type": "info", "text": "just info"}]}]},
    {"start": "second", "screens": [{"id": "first", "body": "1"}, {"id": "second", "body": "2", "next": "$cancel"}]},
]

ODD = [None, True, False, 0, 1, -3, 7, 1.0, 2.5, "", "x", " ", "a b", "$done", "$cancel", [], {}, [1], ["a"], {"a": 1}, "é" * 80, "x" * 65, "y" * 201, "z" * 2001]


def walk(node, path=()):
    yield node, path
    if isinstance(node, dict):
        for key, value in list(node.items()):
            yield from walk(value, path + (key,))
    elif isinstance(node, list):
        for index, value in enumerate(node):
            yield from walk(value, path + (index,))


def at(node, path):
    for step in path:
        node = node[step]
    return node


def mutate(base):
    spec = copy.deepcopy(base)
    for _ in range(rng.choice([1, 1, 1, 1, 2, 3])):
        places = [path for _node, path in walk(spec) if path]
        if not places:
            break
        path = rng.choice(places)
        parent = at(spec, path[:-1])
        key = path[-1]
        action = rng.choice(["replace", "delete", "dupe", "replace"])
        if action == "delete":
            if isinstance(parent, dict):
                del parent[key]
            else:
                parent.pop(key)
        elif action == "dupe" and isinstance(parent, list) and parent:
            parent.append(copy.deepcopy(parent[rng.randrange(len(parent))]))
        else:
            parent[key] = copy.deepcopy(rng.choice(ODD))
    return spec


def run_validate(candidate):
    try:
        app, err = tai.validate_flow_spec(copy.deepcopy(candidate))
    except Exception:  # noqa: BLE001 python crashes are not part of the contract
        return None
    return app, err


class Script(tai.PlexerController):
    kind = "script"

    def __init__(self, replies):
        self.replies = list(replies)
        self.asks = []

    def ask(self, screen, timeout):
        self.asks.append(screen)
        if not self.replies:
            return {"status": "cancelled"}
        return self.replies.pop(0)


def random_replies():
    out = []
    for _ in range(rng.randrange(0, 70)):
        roll = rng.random()
        if roll < 0.04:
            out.append({"status": "cancelled"})
        elif roll < 0.07:
            out.append({"status": "timeout"})
        elif roll < 0.55:
            out.append({"status": "ok", "value": rng.randrange(0, 3)})
        else:
            out.append({"status": "ok", "value": rng.choice(["", "v2", "café", True])})
    return out


def fix_replies(app, replies):
    """Index replies must be valid for the screen they answer; replay with a guard controller."""
    guard = []

    class Guard(tai.PlexerController):
        kind = "guard"

        def __init__(self):
            self.todo = list(replies)

        def ask(self, screen, timeout):
            if not self.todo:
                reply = {"status": "cancelled"}
            else:
                reply = self.todo.pop(0)
            if reply.get("status") == "ok":
                if screen["kind"] == "choice":
                    reply = {"status": "ok", "value": reply["value"] % len(screen["options"]) if isinstance(reply["value"], int) and not isinstance(reply["value"], bool) else 0}
                elif screen["kind"] == "text":
                    reply = {"status": "ok", "value": reply["value"] if isinstance(reply["value"], str) else "t"}
                else:
                    reply = {"status": "ok", "value": True}
            guard.append(reply)
            return reply

    return Guard(), guard


lines = []
cases = 0
seen_specs = []
for base in BASES:
    seen_specs.append(base)
for _ in range(4000):
    seen_specs.append(mutate(rng.choice(BASES)))
for _ in range(60):
    seen_specs.append(rng.choice(ODD))
for candidate in seen_specs:
    result = run_validate(candidate)
    if result is None:
        continue
    app, err = result
    lines.append(json.dumps({"kind": "spec", "spec": candidate, "app": app, "err": err}))
    cases += 1
    if app is not None:
        for _round in range(3):
            controller, used = fix_replies(app, random_replies())
            try:
                outcome = tai.PlexerFlow(copy.deepcopy(app)).run(controller)
            except Exception:  # noqa: BLE001
                continue
            lines.append(json.dumps({"kind": "flow", "app": app, "replies": used, "outcome": outcome}))
            cases += 1
with open(out_path, "w") as handle:
    handle.write("\n".join(lines) + "\n")
print("cases", cases)