LangGraph: agents you can draw

# test_retry_graph.py: run from chapter_18_LangGraph/src/chapters
#   pip install pytest && python -m pytest test_retry_graph.py
import importlib

retry = importlib.import_module("004_Loop_Retry")


def run(outcomes):
    return retry.app.invoke({"test_name": "t", "scripted_outcomes": outcomes,
                             "attempts": 0, "history": [], "verdict": ""})


def test_flaky_passes_on_the_third_attempt():
    out = run(["failed", "failed", "passed"])
    assert out["verdict"] == "FLAKY: passed on attempt 3 after 2 failure(s)"
    assert out["attempts"] == 3


def test_broken_test_stops_at_the_cap():
    out = run(["failed"] * 5)
    assert out["attempts"] == retry.MAX_ATTEMPTS
    assert out["verdict"] == "REAL FAILURE: failed all 3 attempts"


def test_stable_test_runs_once():
    out = run(["passed"])
    assert out["history"] == ["passed"]
    assert out["verdict"] == "STABLE PASS"

import { test, expect } from '@playwright/test';

const URL = 'https://app.thetestingacademy.com/ai/blueprint/learn/langgraph.html';

test('a flaky test pauses for approval, then is quarantined', async ({ page }) => {
  await page.goto(URL);
  await page.getByTestId('lg-scenario').selectOption('flaky');
  await page.getByTestId('lg-run').click();

  await expect(page.getByTestId('lg-node-approve')).toHaveClass(/is-wait/);
  await expect(page.getByTestId('lg-status')).toContainText('Graph paused');

  await page.getByTestId('lg-approve-yes').click();
  await expect(page.getByTestId('lg-verdict'))
    .toHaveText('FLAKY: passed on attempt 3 after 2 failure(s)');
  await expect(page.getByTestId('lg-state')).toContainText('"quarantined": true');
});

test('a broken test stops at the retry cap', async ({ page }) => {
  await page.goto(URL);
  await page.getByTestId('lg-scenario').selectOption('broken');
  await page.getByTestId('lg-max').fill('2');
  await page.getByTestId('lg-max').dispatchEvent('change');
  await page.getByTestId('lg-run').click();

  const attempts = page.getByTestId('lg-log').locator('li[data-kind="bad"]', { hasText: 'run_test' });
  await expect(attempts).toHaveCount(2);
  await expect(page.getByTestId('lg-verdict')).toHaveText('REAL FAILURE: failed all 2 attempts');
});

test('rewinding to a checkpoint restores the earlier state', async ({ page }) => {
  await page.goto(URL);
  await page.getByTestId('lg-run').click();          // flaky scenario, pauses at approve
  await page.getByTestId('lg-checkpoint-2').click();  // right after attempt 1
  await expect(page.getByTestId('lg-state')).toContainText('"attempts": 1');
  await expect(page.getByTestId('lg-node-run_test')).toHaveClass(/is-active/);
});

"""003 - A fork in the road: route a test result to the right next step.

add_conditional_edges(source, router, destinations)
The router is a normal function that reads state and returns the NAME of the
next node. That one function is where an "agent" makes decisions.
"""
from typing import Literal, TypedDict

from langgraph.graph import END, START, StateGraph


class State(TypedDict, total=False):
    test_name: str
    status: str          # passed | failed | flaky
    action: str


def read_result(state: State) -> dict:
    return {}            # in real life: parse a Playwright report here


def route_by_status(state: State) -> Literal["log_pass", "file_bug", "quarantine"]:
    return {"passed": "log_pass", "failed": "file_bug"}.get(state["status"], "quarantine")


def log_pass(state: State) -> dict:
    return {"action": "logged as green"}


def file_bug(state: State) -> dict:
    return {"action": f"Jira bug filed for {state['test_name']}"}


def quarantine(state: State) -> dict:
    return {"action": f"{state['test_name']} moved to quarantine, rerun nightly"}


builder = StateGraph(State)
builder.add_node("read_result", read_result)
builder.add_node("log_pass", log_pass)
builder.add_node("file_bug", file_bug)
builder.add_node("quarantine", quarantine)

builder.add_edge(START, "read_result")
builder.add_conditional_edges("read_result", route_by_status,
                              ["log_pass", "file_bug", "quarantine"])
for leaf in ("log_pass", "file_bug", "quarantine"):
    builder.add_edge(leaf, END)

app = builder.compile()

if __name__ == "__main__":
    for name, status in [("cart_total", "passed"),
                         ("checkout_pay", "failed"),
                         ("login_redirect", "flaky")]:
        out = app.invoke({"test_name": name, "status": status})
        print(f"{name:<15} {status:<7} -> {out['action']}")

    print("\n" + app.get_graph().draw_mermaid())

"""007 - Human in the loop: pause the graph and ask before doing something risky.

interrupt(payload) freezes the graph and hands `payload` to the caller.
You resume it later with Command(resume=answer) on the SAME thread_id.
A checkpointer is required: the pause has to be saved somewhere.

    python 007_Human_In_The_Loop.py        # asks you
    python 007_Human_In_The_Loop.py yes    # answers for you (CI / demos)
"""
import sys
from typing import Literal, TypedDict

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command, interrupt


class State(TypedDict, total=False):
    flaky_tests: list[str]
    approved: bool
    outcome: str


def detect(state: State) -> dict:
    return {"flaky_tests": ["login_redirect", "search_autocomplete"]}


def ask_approval(state: State) -> dict:
    answer = interrupt({
        "question": "Quarantine these tests? (yes/no)",
        "tests": state["flaky_tests"],
    })
    return {"approved": str(answer).strip().lower() in ("y", "yes")}


def route(state: State) -> Literal["quarantine", "skip"]:
    return "quarantine" if state["approved"] else "skip"


def quarantine(state: State) -> dict:
    return {"outcome": f"quarantined {len(state['flaky_tests'])} tests"}


def skip(state: State) -> dict:
    return {"outcome": "left as is, nothing changed"}


builder = StateGraph(State)
builder.add_node("detect", detect)
builder.add_node("ask_approval", ask_approval)
builder.add_node("quarantine", quarantine)
builder.add_node("skip", skip)
builder.add_edge(START, "detect")
builder.add_edge("detect", "ask_approval")
builder.add_conditional_edges("ask_approval", route, ["quarantine", "skip"])
builder.add_edge("quarantine", END)
builder.add_edge("skip", END)

app = builder.compile(checkpointer=InMemorySaver())

if __name__ == "__main__":
    cfg = {"configurable": {"thread_id": "flaky-review-1"}}

    first = app.invoke({}, cfg)
    pause = first["__interrupt__"][0].value
    print("Graph paused. It is asking:")
    print(" ", pause["question"])
    for t in pause["tests"]:
        print("   -", t)
    print("Paused before node:", app.get_state(cfg).next)

    answer = sys.argv[1] if len(sys.argv) > 1 else input("\nYour answer: ")
    final = app.invoke(Command(resume=answer), cfg)
    print("\nResumed with", repr(answer), "->", final["outcome"])