Publish tests/ from marimo

2bef145 verified 4 days ago

12.3 kB

"""Traces record a run, replay it offline, and count how often the choice was right."""

import json

import pytest

from tiny_computer_agent import (
    Action, ActionError, AgentLoop, DecisionMenu, Element, MockBrain, PlannerBrain,
)
from tiny_computer_agent.trace import (
    DONE, NOTHING, TRACE_VERSION, Scenario, Scoreboard, Trace, TraceStep, benchmark, benchmark_traces, load_traces, run_scenario, score, score_replay,
    score_trace,
)


# -- fixtures --------------------------------------------------------------

def frame(*labels):
    return tuple(Element(label, x=i * 100, y=10, width=80, height=24) for i, label in enumerate(labels))


def topbar(*labels):
    return tuple(Element(label, x=i * 100, y=6, width=80, height=20) for i, label in enumerate(labels))


def button(label, x, y=270):
    return Element(label, x=x, y=y, width=90, height=30)


SAVE = Scenario(
    "click 'Save' to save the document",
    frames=[topbar("Files", "Edit", "View"),
            (Element("Name", x=200, y=195, width=200, height=26), button("Save", 200)),
            topbar("Files", "Edit", "View") + (button("Save", 200),)],
    expected="Save", name="save",
)
CANCEL = Scenario(
    "click 'Cancel' to discard the change",
    frames=[(button("Save", 200), button("Cancel", 400))],
    expected="Cancel", name="cancel",
)
DENIED = Scenario(
    "click 'Terminal' to open a shell",
    frames=[topbar("Terminal", "Files")],
    expected="Terminal", name="terminal (denied by policy)",
)


class FirstItemBrain:
    """A baseline: always click the first item in reading order. No model, no idea."""

    name = "first-item"

    def decide(self, goal, observation):
        menu = DecisionMenu(observation)
        if not menu:
            return Action("done", reason="nothing to click")
        return Action("click", label=menu.items[0].label, reason="first item in reading order")


def oracle(scenario):
    return MockBrain([
        {"kind": "click", "label": scenario.expected},
        {"kind": "done", "reason": "clicked the right thing"},
    ])


# -- recording -------------------------------------------------------------

def test_a_trace_records_the_frames_and_the_decision_without_pixels():
    episode = AgentLoop(screen=SAVE.screen(), brain=PlannerBrain()).run(SAVE.goal)
    trace = Trace.record(episode, expected="Save", brain="planner", scenario="save")

    assert trace.expected == "Save" and trace.brain == "planner" and trace.scenario == "save"
    assert trace.frames == len(episode.steps)
    assert trace.choice == "Save"
    # the element index survives; the pixels deliberately do not
    assert trace.steps[-1].elements and all(isinstance(e, Element) for e in trace.steps[-1].elements)
    assert all(step.observation().frame is None for step in trace.steps)


def test_waits_are_not_choices():
    episode = AgentLoop(
        screen=SAVE.screen(),
        brain=MockBrain([{"kind": "wait", "seconds": 0.1}, {"kind": "click", "label": "Save"}]),
    ).run(SAVE.goal)
    trace = Trace.record(episode, expected="Save")
    assert trace.steps[0].action.kind == "wait"
    assert trace.choice == "Save"          # the click, not the wait
    assert trace.decided == "Save"


def test_a_brain_that_says_done_before_choosing_is_recorded_bluntly():
    episode = AgentLoop(
        screen=SAVE.screen(), brain=MockBrain([{"kind": "done", "reason": "assumed"}]),
    ).run(SAVE.goal)
    trace = Trace.record(episode, expected="Save")
    assert trace.choice is None
    assert trace.decided == DONE
    assert score_trace(trace).correct is False


def test_a_trace_that_never_committed_reads_as_nothing():
    trace = Trace(goal="click something")
    assert trace.decided == NOTHING and trace.choice is None
    assert trace.frames == 0


# -- persistence -----------------------------------------------------------

def test_a_trace_round_trips_through_json():
    trace = run_scenario(SAVE, PlannerBrain())
    again = Trace.from_json(trace.to_json())
    assert again.to_dict() == trace.to_dict()
    assert again.to_json() == trace.to_json()
    assert json.loads(trace.to_json())["trace_version"] == TRACE_VERSION


def test_a_trace_saves_and_loads_from_a_directory(tmp_path):
    run_scenario(SAVE, PlannerBrain()).save(tmp_path / "b.json")
    run_scenario(CANCEL, PlannerBrain()).save(tmp_path / "a.json")
    names = [t.name for t in load_traces(tmp_path)]     # directories load in name order
    assert names == ["cancel", "save"]
    assert load_traces(tmp_path / "a.json")[0].name == "cancel"


def test_a_newer_trace_version_is_refused_rather_than_misread():
    with pytest.raises(ValueError, match="newer than this build"):
        Trace.from_dict({"trace_version": TRACE_VERSION + 1, "goal": "g"})


def test_a_corrupt_action_in_a_trace_is_reported_on_load():
    data = {"goal": "g", "steps": [{"index": 0, "action": {"kind": "format_disk"}}]}
    with pytest.raises(ActionError, match="unknown action"):
        Trace.from_dict(data)


def test_unknown_keys_are_ignored_so_a_trace_stays_forward_compatible():
    trace = Trace.from_dict({"goal": "g", "expected": "Save", "extra": 1,
                             "steps": [{"index": 0, "future_field": True}]})
    assert trace.goal == "g" and trace.frames == 1


# -- scoring ---------------------------------------------------------------

def test_score_matches_the_choice_to_the_expected_answer():
    board = benchmark([SAVE, CANCEL], {"planner": lambda s: PlannerBrain()})
    assert board.accuracy == 1.0
    assert all(v.correct for v in board)


def test_score_forgives_ocr_spacing_the_way_the_index_does():
    trace = Trace(goal="g", expected="Save",
                  steps=[TraceStep(0, action=Action("click", label="Save ."))])
    assert score_trace(trace).correct is True


def test_a_wrong_choice_is_scored_with_the_alternative_it_picked():
    verdict = score(SAVE, FirstItemBrain())
    assert verdict.correct is False
    assert verdict.chosen == "Files"
    assert "not 'Save'" in verdict.detail


def test_a_trace_with_no_expectation_cannot_be_scored_and_says_so():
    verdict = score_trace(run_scenario(SAVE, PlannerBrain()))
    # run_scenario always records ground truth, so drop it to check the guard
    unjudged = score_trace(Trace(goal="g", steps=[TraceStep(0, action=Action("click", label="Save"))]))
    assert verdict.correct is True
    assert unjudged.correct is False and "no expected answer" in unjudged.detail


def test_a_refused_decision_is_scored_on_the_choice_but_flagged():
    """A denied click is still what the brain chose - and it must not look clean."""
    verdict = score(DENIED, FirstItemBrain())
    assert verdict.chosen == "Terminal" and verdict.correct is True
    assert verdict.blocked is True
    assert "refused" in verdict.line()


def test_the_scoreboard_counts_blocked_decisions_separately_from_accuracy():
    board = benchmark([SAVE, DENIED], {"first-item": lambda s: FirstItemBrain()})
    assert board.total == 2 and board.blocked == 1
    assert "stopped by the harness" in board.render()


# -- replay ----------------------------------------------------------------

def test_replay_answers_the_recorded_frames_with_a_different_brain():
    recorded = run_scenario(SAVE, PlannerBrain())
    assert recorded.choice == "Save"
    replayed = recorded.replay(FirstItemBrain())
    assert replayed.brain == "first-item"
    assert replayed.choice == "Files"                       # same frames, a worse brain
    assert replayed.frames == recorded.frames
    assert replayed.steps[0].elements == recorded.steps[0].elements


def test_replay_keeps_the_ground_truth_so_the_replay_can_be_scored():
    recorded = run_scenario(CANCEL, FirstItemBrain())
    replayed = recorded.replay(PlannerBrain())
    assert replayed.expected == "Cancel"
    assert score_trace(replayed).correct is True


def test_a_brain_that_explodes_during_replay_is_a_scored_error_not_a_crash():
    class Exploding:
        name = "exploding"

        def decide(self, goal, observation):
            raise RuntimeError("model died")

    replayed = run_scenario(SAVE, PlannerBrain()).replay(Exploding())
    assert replayed.decided == NOTHING
    assert "RuntimeError: model died" in replayed.steps[0].error
    assert score_trace(replayed).correct is False


def test_score_replay_is_the_single_trace_form_of_the_roadmap_step():
    recorded = run_scenario(SAVE, PlannerBrain())
    assert score_replay(recorded, PlannerBrain()).correct is True
    assert score_replay(recorded, FirstItemBrain()).correct is False


def test_benchmark_traces_replays_a_recorded_suite_against_one_brain():
    traces = [run_scenario(SAVE, PlannerBrain())]        # a trace of a run that got it right
    assert benchmark_traces(traces, PlannerBrain()).accuracy == 1.0
    # the same frames, a worse brain: the trace is the fixed part, the brain is the variable
    replayed = benchmark_traces(traces, FirstItemBrain())
    assert replayed.brains() == ["first-item"]
    assert replayed.accuracy == 0.0


# -- the suite, and its table ---------------------------------------------

def test_benchmark_scores_every_brain_on_every_scenario_with_a_fresh_brain_each_time():
    seen = []

    def make(scenario):
        seen.append(scenario.name)
        return PlannerBrain()

    board = benchmark([SAVE, CANCEL], {"planner": make, "oracle": oracle, "first-item": lambda s: FirstItemBrain()})
    assert board.total == 6
    assert board.brains() == ["planner", "oracle", "first-item"]
    assert board.by_brain() == {"planner": 1.0, "oracle": 1.0, "first-item": 0.0}
    assert seen == ["save", "cancel"]                 # a fresh brain per scenario
    assert {row["correct"] for row in board.to_rows()} == {True, False}


def test_the_scoreboard_table_has_the_columns_a_notebook_wants():
    board = benchmark([SAVE], {"planner": lambda s: PlannerBrain()})
    assert board.to_rows() == [
        {"brain": "planner", "scenario": "save", "expected": "Save", "chose": "Save", "correct": True}
    ]
    assert board.confusion() == {("Save", "Save"): 1}


def test_an_empty_scoreboard_is_zero_not_a_division_error():
    board = Scoreboard()
    assert board.total == 0 and board.accuracy == 0.0 and board.brains() == []
    assert board.to_rows() == [] and "0/0" in board.render()


# -- scenarios -------------------------------------------------------------

def test_a_scenario_needs_at_least_one_frame():
    with pytest.raises(ValueError, match="at least one frame"):
        Scenario("g", frames=[], expected="Save")


def test_a_scenario_hands_out_a_fresh_screen_every_time():
    scenario = Scenario("g", frames=[frame("Save")], expected="Save")
    first, second = scenario.screen(), scenario.screen()
    assert first is not second
    assert first.captures == second.captures == 0
    assert scenario.frames[0][0].label == "Save"


def test_render_is_a_replayable_summary_of_the_trace():
    trace = run_scenario(DENIED, FirstItemBrain())
    text = trace.render()
    assert "trace: terminal (denied by policy)" in text
    assert "brain: first-item" in text
    assert "expected: 'Terminal'" in text and "chose: 'Terminal'" in text
    assert "the harness refused it" in text
    assert "click" in text


# -- the command line ------------------------------------------------------

def test_cli_records_a_trace_and_scores_it(tmp_path, capsys):
    from tiny_computer_agent import __main__ as cli

    target = tmp_path / "files.json"
    assert cli.main(["--script", '[{"kind": "click", "label": "Files"}]',
                     "--trace", str(target), "--expected", "Files"]) == 0
    out = capsys.readouterr().out
    assert "trace: wrote" in out and "2 frames" in out
    assert "correct" in out

    trace = Trace.load(target)
    assert trace.expected == "Files" and trace.choice == "Files"
    assert score_trace(trace).correct is True


def test_cli_writes_a_trace_even_without_an_expected_answer(tmp_path, capsys):
    from tiny_computer_agent import __main__ as cli

    target = tmp_path / "unscored.json"
    assert cli.main(["--script", '[{"kind": "click", "label": "Files"}]',
                     "--trace", str(target)]) == 0
    assert "score:" not in capsys.readouterr().out
    trace = Trace.load(target)
    assert trace.expected is None and score_trace(trace).correct is False

	"""Traces record a run, replay it offline, and count how often the choice was right."""

	import json

	import pytest

	from tiny_computer_agent import (
	Action, ActionError, AgentLoop, DecisionMenu, Element, MockBrain, PlannerBrain,
	)
	from tiny_computer_agent.trace import (
	DONE, NOTHING, TRACE_VERSION, Scenario, Scoreboard, Trace, TraceStep, benchmark, benchmark_traces, load_traces, run_scenario, score, score_replay,
	score_trace,
	)


	# -- fixtures --------------------------------------------------------------

	def frame(*labels):
	return tuple(Element(label, x=i * 100, y=10, width=80, height=24) for i, label in enumerate(labels))


	def topbar(*labels):
	return tuple(Element(label, x=i * 100, y=6, width=80, height=20) for i, label in enumerate(labels))


	def button(label, x, y=270):
	return Element(label, x=x, y=y, width=90, height=30)


	SAVE = Scenario(
	"click 'Save' to save the document",
	frames=[topbar("Files", "Edit", "View"),
	(Element("Name", x=200, y=195, width=200, height=26), button("Save", 200)),
	topbar("Files", "Edit", "View") + (button("Save", 200),)],
	expected="Save", name="save",
	)
	CANCEL = Scenario(
	"click 'Cancel' to discard the change",
	frames=[(button("Save", 200), button("Cancel", 400))],
	expected="Cancel", name="cancel",
	)
	DENIED = Scenario(
	"click 'Terminal' to open a shell",
	frames=[topbar("Terminal", "Files")],
	expected="Terminal", name="terminal (denied by policy)",
	)


	class FirstItemBrain:
	"""A baseline: always click the first item in reading order. No model, no idea."""

	name = "first-item"

	def decide(self, goal, observation):
	menu = DecisionMenu(observation)
	if not menu:
	return Action("done", reason="nothing to click")
	return Action("click", label=menu.items[0].label, reason="first item in reading order")


	def oracle(scenario):
	return MockBrain([
	{"kind": "click", "label": scenario.expected},
	{"kind": "done", "reason": "clicked the right thing"},
	])


	# -- recording -------------------------------------------------------------

	def test_a_trace_records_the_frames_and_the_decision_without_pixels():
	episode = AgentLoop(screen=SAVE.screen(), brain=PlannerBrain()).run(SAVE.goal)
	trace = Trace.record(episode, expected="Save", brain="planner", scenario="save")

	assert trace.expected == "Save" and trace.brain == "planner" and trace.scenario == "save"
	assert trace.frames == len(episode.steps)
	assert trace.choice == "Save"
	# the element index survives; the pixels deliberately do not
	assert trace.steps[-1].elements and all(isinstance(e, Element) for e in trace.steps[-1].elements)
	assert all(step.observation().frame is None for step in trace.steps)


	def test_waits_are_not_choices():
	episode = AgentLoop(
	screen=SAVE.screen(),
	brain=MockBrain([{"kind": "wait", "seconds": 0.1}, {"kind": "click", "label": "Save"}]),
	).run(SAVE.goal)
	trace = Trace.record(episode, expected="Save")
	assert trace.steps[0].action.kind == "wait"
	assert trace.choice == "Save" # the click, not the wait
	assert trace.decided == "Save"


	def test_a_brain_that_says_done_before_choosing_is_recorded_bluntly():
	episode = AgentLoop(
	screen=SAVE.screen(), brain=MockBrain([{"kind": "done", "reason": "assumed"}]),
	).run(SAVE.goal)
	trace = Trace.record(episode, expected="Save")
	assert trace.choice is None
	assert trace.decided == DONE
	assert score_trace(trace).correct is False


	def test_a_trace_that_never_committed_reads_as_nothing():
	trace = Trace(goal="click something")
	assert trace.decided == NOTHING and trace.choice is None
	assert trace.frames == 0


	# -- persistence -----------------------------------------------------------

	def test_a_trace_round_trips_through_json():
	trace = run_scenario(SAVE, PlannerBrain())
	again = Trace.from_json(trace.to_json())
	assert again.to_dict() == trace.to_dict()
	assert again.to_json() == trace.to_json()
	assert json.loads(trace.to_json())["trace_version"] == TRACE_VERSION


	def test_a_trace_saves_and_loads_from_a_directory(tmp_path):
	run_scenario(SAVE, PlannerBrain()).save(tmp_path / "b.json")
	run_scenario(CANCEL, PlannerBrain()).save(tmp_path / "a.json")
	names = [t.name for t in load_traces(tmp_path)] # directories load in name order
	assert names == ["cancel", "save"]
	assert load_traces(tmp_path / "a.json")[0].name == "cancel"


	def test_a_newer_trace_version_is_refused_rather_than_misread():
	with pytest.raises(ValueError, match="newer than this build"):
	Trace.from_dict({"trace_version": TRACE_VERSION + 1, "goal": "g"})


	def test_a_corrupt_action_in_a_trace_is_reported_on_load():
	data = {"goal": "g", "steps": [{"index": 0, "action": {"kind": "format_disk"}}]}
	with pytest.raises(ActionError, match="unknown action"):
	Trace.from_dict(data)


	def test_unknown_keys_are_ignored_so_a_trace_stays_forward_compatible():
	trace = Trace.from_dict({"goal": "g", "expected": "Save", "extra": 1,
	"steps": [{"index": 0, "future_field": True}]})
	assert trace.goal == "g" and trace.frames == 1


	# -- scoring ---------------------------------------------------------------

	def test_score_matches_the_choice_to_the_expected_answer():
	board = benchmark([SAVE, CANCEL], {"planner": lambda s: PlannerBrain()})
	assert board.accuracy == 1.0
	assert all(v.correct for v in board)


	def test_score_forgives_ocr_spacing_the_way_the_index_does():
	trace = Trace(goal="g", expected="Save",
	steps=[TraceStep(0, action=Action("click", label="Save ."))])
	assert score_trace(trace).correct is True


	def test_a_wrong_choice_is_scored_with_the_alternative_it_picked():
	verdict = score(SAVE, FirstItemBrain())
	assert verdict.correct is False
	assert verdict.chosen == "Files"
	assert "not 'Save'" in verdict.detail


	def test_a_trace_with_no_expectation_cannot_be_scored_and_says_so():
	verdict = score_trace(run_scenario(SAVE, PlannerBrain()))
	# run_scenario always records ground truth, so drop it to check the guard
	unjudged = score_trace(Trace(goal="g", steps=[TraceStep(0, action=Action("click", label="Save"))]))
	assert verdict.correct is True
	assert unjudged.correct is False and "no expected answer" in unjudged.detail


	def test_a_refused_decision_is_scored_on_the_choice_but_flagged():
	"""A denied click is still what the brain chose - and it must not look clean."""
	verdict = score(DENIED, FirstItemBrain())
	assert verdict.chosen == "Terminal" and verdict.correct is True
	assert verdict.blocked is True
	assert "refused" in verdict.line()


	def test_the_scoreboard_counts_blocked_decisions_separately_from_accuracy():
	board = benchmark([SAVE, DENIED], {"first-item": lambda s: FirstItemBrain()})
	assert board.total == 2 and board.blocked == 1
	assert "stopped by the harness" in board.render()


	# -- replay ----------------------------------------------------------------

	def test_replay_answers_the_recorded_frames_with_a_different_brain():
	recorded = run_scenario(SAVE, PlannerBrain())
	assert recorded.choice == "Save"
	replayed = recorded.replay(FirstItemBrain())
	assert replayed.brain == "first-item"
	assert replayed.choice == "Files" # same frames, a worse brain
	assert replayed.frames == recorded.frames
	assert replayed.steps[0].elements == recorded.steps[0].elements


	def test_replay_keeps_the_ground_truth_so_the_replay_can_be_scored():
	recorded = run_scenario(CANCEL, FirstItemBrain())
	replayed = recorded.replay(PlannerBrain())
	assert replayed.expected == "Cancel"
	assert score_trace(replayed).correct is True


	def test_a_brain_that_explodes_during_replay_is_a_scored_error_not_a_crash():
	class Exploding:
	name = "exploding"

	def decide(self, goal, observation):
	raise RuntimeError("model died")

	replayed = run_scenario(SAVE, PlannerBrain()).replay(Exploding())
	assert replayed.decided == NOTHING
	assert "RuntimeError: model died" in replayed.steps[0].error
	assert score_trace(replayed).correct is False


	def test_score_replay_is_the_single_trace_form_of_the_roadmap_step():
	recorded = run_scenario(SAVE, PlannerBrain())
	assert score_replay(recorded, PlannerBrain()).correct is True
	assert score_replay(recorded, FirstItemBrain()).correct is False


	def test_benchmark_traces_replays_a_recorded_suite_against_one_brain():
	traces = [run_scenario(SAVE, PlannerBrain())] # a trace of a run that got it right
	assert benchmark_traces(traces, PlannerBrain()).accuracy == 1.0
	# the same frames, a worse brain: the trace is the fixed part, the brain is the variable
	replayed = benchmark_traces(traces, FirstItemBrain())
	assert replayed.brains() == ["first-item"]
	assert replayed.accuracy == 0.0


	# -- the suite, and its table ---------------------------------------------

	def test_benchmark_scores_every_brain_on_every_scenario_with_a_fresh_brain_each_time():
	seen = []

	def make(scenario):
	seen.append(scenario.name)
	return PlannerBrain()

	board = benchmark([SAVE, CANCEL], {"planner": make, "oracle": oracle, "first-item": lambda s: FirstItemBrain()})
	assert board.total == 6
	assert board.brains() == ["planner", "oracle", "first-item"]
	assert board.by_brain() == {"planner": 1.0, "oracle": 1.0, "first-item": 0.0}
	assert seen == ["save", "cancel"] # a fresh brain per scenario
	assert {row["correct"] for row in board.to_rows()} == {True, False}


	def test_the_scoreboard_table_has_the_columns_a_notebook_wants():
	board = benchmark([SAVE], {"planner": lambda s: PlannerBrain()})
	assert board.to_rows() == [
	{"brain": "planner", "scenario": "save", "expected": "Save", "chose": "Save", "correct": True}
	]
	assert board.confusion() == {("Save", "Save"): 1}


	def test_an_empty_scoreboard_is_zero_not_a_division_error():
	board = Scoreboard()
	assert board.total == 0 and board.accuracy == 0.0 and board.brains() == []
	assert board.to_rows() == [] and "0/0" in board.render()


	# -- scenarios -------------------------------------------------------------

	def test_a_scenario_needs_at_least_one_frame():
	with pytest.raises(ValueError, match="at least one frame"):
	Scenario("g", frames=[], expected="Save")


	def test_a_scenario_hands_out_a_fresh_screen_every_time():
	scenario = Scenario("g", frames=[frame("Save")], expected="Save")
	first, second = scenario.screen(), scenario.screen()
	assert first is not second
	assert first.captures == second.captures == 0
	assert scenario.frames[0][0].label == "Save"


	def test_render_is_a_replayable_summary_of_the_trace():
	trace = run_scenario(DENIED, FirstItemBrain())
	text = trace.render()
	assert "trace: terminal (denied by policy)" in text
	assert "brain: first-item" in text
	assert "expected: 'Terminal'" in text and "chose: 'Terminal'" in text
	assert "the harness refused it" in text
	assert "click" in text


	# -- the command line ------------------------------------------------------

	def test_cli_records_a_trace_and_scores_it(tmp_path, capsys):
	from tiny_computer_agent import __main__ as cli

	target = tmp_path / "files.json"
	assert cli.main(["--script", '[{"kind": "click", "label": "Files"}]',
	"--trace", str(target), "--expected", "Files"]) == 0
	out = capsys.readouterr().out
	assert "trace: wrote" in out and "2 frames" in out
	assert "correct" in out

	trace = Trace.load(target)
	assert trace.expected == "Files" and trace.choice == "Files"
	assert score_trace(trace).correct is True


	def test_cli_writes_a_trace_even_without_an_expected_answer(tmp_path, capsys):
	from tiny_computer_agent import __main__ as cli

	target = tmp_path / "unscored.json"
	assert cli.main(["--script", '[{"kind": "click", "label": "Files"}]',
	"--trace", str(target)]) == 0
	assert "score:" not in capsys.readouterr().out
	trace = Trace.load(target)
	assert trace.expected is None and score_trace(trace).correct is False