Publish tests/ from marimo

tests/test_brain.py

@@ -0,0 +1,155 @@
+import pytest
+
+from tiny_computer_agent import (
+    Action, BrainError, DecisionMenu, Element, GoalUnreachable, ModelBrain,
+    Observation, extract_json_object, extract_json_objects, parse_decision,
+)
+
+FRAME = Observation(elements=(
+    Element("Files", x=10, y=5, width=40, height=20),
+    Element("Save", x=200, y=260, width=60, height=20),
+    Element("Cancel", x=300, y=260, width=60, height=20),
+), text="untitled document")
+
+
+def test_menu_is_numbered_in_reading_order_and_capped():
+    observation = Observation(elements=tuple(
+        Element(f"b{i}", x=(i % 2) * 100, y=(i // 2) * 50) for i in range(6)
+    ))
+    menu = DecisionMenu(observation, max_items=3)
+    assert [item.index for item in menu] == [1, 2, 3]
+    assert len(menu) == 3 and bool(menu)
+    assert "1." in menu.render()
+
+
+def test_menu_lookup_tolerates_what_models_send():
+    menu = DecisionMenu(FRAME)
+    assert menu.item(1).label == "Files"
+    assert menu.item("2").label == "Save"
+    assert menu.item(2.0).label == "Save"
+    assert menu.item(0) is None and menu.item(99) is None and menu.item("x") is None
+
+
+def test_menu_resolves_labels_through_normalisation():
+    menu = DecisionMenu(Observation(elements=(Element("Save ."),)))
+    assert menu.find("save").label == "Save ."
+    assert menu.find("nothing").__class__ is type(None)
+
+
+def test_json_objects_are_extracted_in_order_and_the_last_wins():
+    raw = 'First {"kind": "click", "index": 1} then {"kind": "done", "reason": "ok"}'
+    assert [o["kind"] for o in extract_json_objects(raw)] == ["click", "done"]
+    # a model that echoes its prompt puts its own answer last
+    assert extract_json_object(raw)["kind"] == "done"
+    with pytest.raises(BrainError, match="no JSON object"):
+        extract_json_object("no json here")
+
+
+def test_parse_decision_json_index():
+    action = parse_decision('{"kind": "click", "index": 2, "reason": "saves"}', FRAME)
+    assert action.kind == "click" and action.label == "Save"
+
+
+def test_parse_decision_echo_of_the_prompt_is_rejected():
+    echoed = '{"kind": "click", "index": 1} {"kind": "click", "index": 2}'
+    assert parse_decision(echoed, FRAME).label == "Save"          # last object wins
+    with pytest.raises(BrainError, match="echo"):
+        parse_decision(DecisionMenu(FRAME).render(), FRAME)        # 3+ labels mentioned
+
+
+def test_parse_decision_bare_number_and_prose_label():
+    assert parse_decision("2", FRAME).label == "Save"
+    assert parse_decision('Click on "Save" in the dialogue.', FRAME).label == "Save"
+    assert parse_decision("The goal is already done.", FRAME).kind == "done"
+
+
+def test_parse_decision_rejects_an_index_outside_the_menu():
+    with pytest.raises(BrainError, match="not in the"):
+        parse_decision('{"kind": "click", "index": 99}', FRAME)
+
+
+def test_parse_decision_rejects_a_label_that_is_not_on_screen():
+    with pytest.raises(BrainError, match="not on screen"):
+        parse_decision('{"kind": "click", "label": "Format disk"}', FRAME)
+
+
+def test_parse_decision_tolerates_a_garbled_verb_when_a_label_is_present():
+    # 256M models really do emit {"kind": "file", "label": "Save"}
+    action = parse_decision('{"kind": "file", "label": "Save"}', FRAME)
+    assert action.kind == "click" and action.label == "Save"
+
+
+def test_parse_decision_empty_and_useless_replies():
+    for raw in ["", "   ", "I am not sure what to do, sorry."]:
+        with pytest.raises(BrainError):
+            parse_decision(raw, FRAME)
+
+
+class _ScriptedModel(ModelBrain):
+    name = "scripted"
+
+    def __init__(self, replies, **kwargs):
+        super().__init__("fake/model", **kwargs)
+        self._replies = list(replies)
+        self.prompts = []
+
+    def _generate(self, prompt, observation):
+        self.prompts.append(prompt)
+        return self._replies.pop(0)
+
+
+def test_model_brain_tells_the_model_what_it_already_tried():
+    brain = _ScriptedModel(['{"kind": "click", "index": 2}', '{"kind": "done", "reason": "ok"}'])
+    first = brain.decide("save", FRAME)
+    assert first.label == "Save"
+    assert "Already attempted" not in brain.prompts[0]  # nothing tried yet
+
+    second = brain.decide("save", FRAME)                 # same screen -> no progress
+    prompt = brain.prompts[1]
+    assert "Already attempted" in prompt
+    assert "screen did NOT change" in prompt
+    assert "Choose a DIFFERENT action" in prompt
+    assert second.kind == "done"
+
+
+def test_model_brain_history_is_bounded():
+    brain = _ScriptedModel(['{"kind": "wait", "seconds": 1}'] * 6)
+    for _ in range(6):
+        brain.decide("wait", FRAME)
+    assert len(brain.history) == 3
+
+
+def test_prompt_carries_the_goal_the_menu_and_the_done_option():
+    brain = _ScriptedModel(['{"kind": "done"}'])
+    brain.decide("save the document", FRAME)
+    prompt = brain.prompts[0]
+    assert "save the document" in prompt
+    assert "2. Save  (at" in prompt
+    assert '"kind": "done"' in prompt
+    assert "<number from the list>" in prompt      # placeholders, not worked examples
+
+
+def test_text_only_brains_stop_when_there_is_nothing_to_click():
+    class NeedsIndex(_ScriptedModel):
+        requires_index = True
+
+    with pytest.raises(GoalUnreachable, match="no clickable elements"):
+        NeedsIndex(['{"kind": "done"}']).decide("save", Observation(elements=()))
+
+
+def test_a_right_answer_in_the_wrong_field_is_recovered():
+    """Real 0.5B output: named the correct element but put it in "reason"."""
+    action = parse_decision('{"kind": "click", "index": 0, "reason": "Save"}', FRAME)
+    assert action.kind == "click" and action.label == "Save"
+    assert "recovered" in action.reason
+
+
+def test_recovery_still_refuses_an_unresolvable_index():
+    with pytest.raises(BrainError, match="not in the"):
+        parse_decision('{"kind": "click", "index": 42, "reason": "no idea"}', FRAME)
+
+
+def test_recovery_is_not_ambiguous():
+    # two labels named and no usable index -> refuse rather than guess
+    with pytest.raises(BrainError):
+        parse_decision('{"kind": "click", "index": 0, "reason": "Save or Cancel"}', FRAME)

tests/test_cli.py

@@ -31,3 +31,42 @@ def test_failure_exit_code(capsys):
     # the planner gives up when the target never appears, so the CLI exits non-zero
     assert cli.main(["click 'Nonexistent button'"]) == 1
     assert "gave up" in capsys.readouterr().out
+
+
+def test_doctor_exits_with_a_status(capsys):
+    assert cli.main(["--doctor"]) in (0, 1)
+    assert "preflight" in capsys.readouterr().out
+
+
+def test_live_is_refused_without_a_display(capsys, monkeypatch):
+    from tiny_computer_agent.desktop import display_available
+
+    if display_available()[0]:
+        pytest.skip("this machine has a display, so --live would really run")
+    assert cli.main(["--live"]) == 2
+    captured = capsys.readouterr()
+    assert "refusing to run live" in captured.err
+    assert "no DISPLAY" in captured.out
+
+
+def test_screenshot_replay_needs_no_display(capsys):
+    """The whole pipeline - OCR index, menu, action - on a still image."""
+    from tiny_computer_agent.perception import detector_available
+
+    if not detector_available()[0]:
+        pytest.skip("no OCR engine installed")
+    assert cli.main(["--image", "example", "--brain", "planner"]) == 0
+    out = capsys.readouterr().out
+    assert "screenshot replay" in out
+    assert "elements detected" in out
+    assert "would click" in out
+
+
+def test_a_missing_image_is_reported_not_crashed(capsys):
+    assert cli.main(["--image", "/nonexistent/shot.png"]) == 2
+    assert "could not start" in capsys.readouterr().err
+
+
+def test_vlm_brain_requires_a_live_screen(capsys):
+    assert cli.main(["--brain", "vlm"]) == 2
+    assert "needs --live" in capsys.readouterr().err

tests/test_desktop.py

@@ -0,0 +1,94 @@
+import pytest
+
+import platform
+
+from tiny_computer_agent import (
+    AgentLoop, DryRunExecutor, Element, FakeScreen, MockBrain, PlannerBrain,
+    PyAutoGuiExecutor,
+)
+from tiny_computer_agent.desktop import (
+    DEFAULT_KILL_SWITCH, Check, build_agent, display_available, doctor, run_checks,
+)
+
+FRAME = FakeScreen([(Element("Save"),)])
+
+
+class StubDetector:
+    name = "stub"
+
+    def detect(self, image):
+        return ()
+
+
+def test_display_detection_follows_the_environment(monkeypatch):
+    monkeypatch.delenv("DISPLAY", raising=False)
+    monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
+    ok, detail = display_available()
+    if platform.system() in ("Linux", "FreeBSD"):
+        assert ok is False and "DISPLAY" in detail
+    monkeypatch.setenv("DISPLAY", ":0")
+    assert display_available()[0] is True
+
+
+def test_checks_cover_every_piece_the_harness_needs():
+    names = [c.name for c in run_checks(brain="planner", detector="ocr")]
+    for expected in ("display server", "pillow", "element detector (ocr)",
+                     "execution mode", "kill switch"):
+        assert expected in names
+
+
+def test_pyautogui_is_only_required_for_real_input():
+    dry = {c.name: c for c in run_checks(real_input=False)}
+    live = {c.name: c for c in run_checks(real_input=True)}
+    assert dry["input backend (pyautogui)"].required is False
+    assert live["input backend (pyautogui)"].required is True
+    assert live["execution mode"].detail.startswith("REAL")
+
+
+def test_the_planner_needs_no_model_runtime():
+    checks = {c.name: c for c in run_checks(brain="planner")}
+    assert checks["model runtime"].ok
+    assert "no weights" in checks["model runtime"].detail
+
+
+def test_a_model_brain_reports_its_runtime():
+    checks = {c.name: c for c in run_checks(brain="text-planner")}
+    assert "model runtime" in checks
+    assert "transformers" in checks["model runtime"].detail
+
+
+def test_doctor_gives_a_verdict_and_a_fix(monkeypatch):
+    monkeypatch.delenv("DISPLAY", raising=False)
+    monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
+    report = doctor(brain="planner")
+    assert "preflight" in report
+    assert "fix:" in report or "READY" in report
+
+
+def test_a_check_renders_its_fix_hint():
+    rendered = Check("display server", False, "no display", fix="run on the desktop").render()
+    assert "[FAIL]" in rendered and "fix: run on the desktop" in rendered
+
+
+def test_build_agent_defaults_to_dry_run_and_the_planner():
+    loop = build_agent(screen=FRAME, detector=StubDetector(), brain="planner")
+    assert isinstance(loop, AgentLoop)
+    assert isinstance(loop.brain, PlannerBrain)
+    assert isinstance(loop.executor, DryRunExecutor)
+    assert loop.policy.dry_run is True
+    assert loop.policy.step_budget == 10
+
+
+def test_build_agent_real_input_swaps_the_executor():
+    loop = build_agent(screen=FRAME, detector=StubDetector(), brain=MockBrain([]),
+                       dry_run=False, max_steps=3)
+    assert isinstance(loop.executor, PyAutoGuiExecutor)
+    assert loop.policy.dry_run is False
+
+
+def test_build_agent_wires_the_kill_switch(tmp_path):
+    marker = tmp_path / "STOP"
+    loop = build_agent(screen=FRAME, detector=StubDetector(), brain=MockBrain([]),
+                       kill_switch_path=str(marker))
+    assert loop.policy.kill_switch_path == str(marker)
+    assert DEFAULT_KILL_SWITCH.endswith("tiny-computer-agent.STOP")

tests/test_grid_fallback.py

@@ -0,0 +1,86 @@
+"""The grid fallback: an action space that needs nothing, and its guard rails."""
+
+import pytest
+from PIL import Image
+
+from tiny_computer_agent import Action, Element, Observation, Policy, PolicyViolation
+from tiny_computer_agent.perception import (
+    AutoDetector, GridDetector, OcrDetector, detector_available, make_detector,
+)
+
+CELL = Element("cell 12", kind="grid", x=600, y=0, width=60, height=57)
+
+
+class FixedOcr:
+    name = "fixed"
+
+    def __init__(self, elements=()):
+        self.elements = tuple(elements)
+
+    def detect(self, image):
+        return self.elements
+
+
+def test_grid_detector_needs_nothing_at_all():
+    elements = GridDetector(rows=2, cols=2).detect(Image.new("RGB", (100, 80)))
+    assert [e.label for e in elements] == ["cell 1", "cell 2", "cell 3", "cell 4"]
+    assert all(e.kind == "grid" for e in elements)
+    assert elements[0].center == (25, 20)
+    assert elements[3].center == (75, 60)
+
+
+def test_grid_detector_rejects_a_useless_shape():
+    with pytest.raises(ValueError):
+        GridDetector(rows=0)
+
+
+def test_grid_detector_is_always_available_and_default_off():
+    assert detector_available("grid")[0] is True
+    assert isinstance(make_detector("grid"), GridDetector)
+
+
+def test_auto_uses_ocr_when_it_finds_text():
+    detector = AutoDetector(ocr=FixedOcr((Element("Save"),)), grid=GridDetector(rows=1, cols=1))
+    assert [e.label for e in detector.detect(Image.new("RGB", (10, 10)))] == ["Save"]
+
+
+def test_auto_falls_back_to_the_grid_when_ocr_finds_nothing():
+    detector = AutoDetector(ocr=FixedOcr(()), grid=GridDetector(rows=1, cols=2))
+    assert [e.label for e in detector.detect(Image.new("RGB", (10, 10)))] == ["cell 1", "cell 2"]
+
+
+def test_a_grid_click_is_fine_while_planning():
+    Policy(dry_run=True).check(Action("click", label="cell 12"), Observation(elements=(CELL,)))
+
+
+def test_a_grid_click_is_refused_with_real_input():
+    """A blind click on a real desktop must be a deliberate choice, not a default."""
+    with pytest.raises(PolicyViolation, match="blind click"):
+        Policy(dry_run=False).check(Action("click", label="cell 12"), Observation(elements=(CELL,)))
+
+
+def test_a_grid_click_is_refused_even_when_the_index_is_stale():
+    with pytest.raises(PolicyViolation, match="blind click"):
+        Policy(dry_run=False).check(Action("click", label="cell 47"), Observation())
+
+
+def test_a_grid_click_is_allowed_when_asked_for():
+    Policy(dry_run=False, allow_grid_clicks=True).check(
+        Action("click", label="cell 12"), Observation(elements=(CELL,))
+    )
+
+
+def test_an_ordinary_click_is_unaffected_by_the_grid_guard():
+    Policy(dry_run=False).check(
+        Action("click", label="Save"),
+        Observation(elements=(Element("Save", x=1, y=2, width=3, height=4),)),
+    )
+
+
+@pytest.mark.skipif(not detector_available("ocr")[0], reason="no OCR engine installed")
+def test_real_ocr_beats_the_grid_on_a_real_screenshot():
+    from tiny_computer_agent.demo import fake_desktop
+
+    elements = AutoDetector().detect(fake_desktop())
+    assert elements and any(e.kind != "grid" for e in elements)
+    assert isinstance(OcrDetector(), OcrDetector)

tests/test_loop.py

@@ -46,7 +46,8 @@ def test_screen_is_reread_after_every_executed_action():
 
 def test_step_budget_is_enforced():
     screen = FakeScreen([(Element("Save"),)])
-    brain = MockBrain([{"kind": "click", "label": "Save"}] * 10)
+    # distinct actions, so the budget is what stops it rather than the no-progress guard
+    brain = MockBrain([{"kind": "wait", "seconds": 0.5 + i} for i in range(10)])
     episode = AgentLoop(
         screen=screen, brain=brain, policy=Policy(step_budget=3)
     ).run("click Save")
@@ -56,6 +57,52 @@ def test_step_budget_is_enforced():
     assert "budget of 3" in episode.stop_reason
 
 
+def test_a_flailing_agent_is_stopped_instead_of_burning_the_budget():
+    """Repeating one action against an unchanged screen is stuck, not persistent."""
+    brain = MockBrain([{"kind": "click", "label": "Save"}] * 10)
+    episode = AgentLoop(
+        screen=FakeScreen([(Element("Save"),)]), brain=brain,
+        policy=Policy(step_budget=10),
+    ).run("click Save")
+
+    assert episode.success is False
+    assert "no progress" in episode.stop_reason
+    assert len(episode.steps) < 10
+    assert "unchanged screen" in episode.steps[-1].error
+
+
+def test_progress_resets_the_repeat_counter():
+    """Two different actions on an unchanged screen is not flailing."""
+    brain = MockBrain([
+        {"kind": "wait", "seconds": 0.5},
+        {"kind": "click", "label": "Save"},
+        {"kind": "wait", "seconds": 0.5},
+        {"kind": "click", "label": "Save"},
+        {"kind": "done", "reason": "assumed saved"},
+    ])
+    episode = AgentLoop(
+        screen=FakeScreen([(Element("Save"),)]), brain=brain,
+        policy=Policy(step_budget=10),
+    ).run("save")
+
+    assert episode.success
+    assert len(episode.steps) == 5
+
+
+def test_a_dead_screen_is_reported_as_an_episode_not_a_traceback():
+    class Blind:
+        name = "blind"
+
+        def capture(self):
+            raise RuntimeError("X connection failed")
+
+    episode = AgentLoop(screen=Blind(), brain=MockBrain([{"kind": "done"}])).run("anything")
+
+    assert episode.success is False
+    assert episode.stop_reason.startswith("capture failed: RuntimeError")
+    assert episode.steps == []
+
+
 def test_policy_refusal_stops_the_run_and_nothing_executes():
     executor = DryRunExecutor()
     episode = AgentLoop(

tests/test_perception.py

@@ -0,0 +1,119 @@
+import pytest
+
+from tiny_computer_agent import Element, FakeScreen, Observation
+from tiny_computer_agent.demo import fake_desktop, write_fake_desktop
+from tiny_computer_agent.observe import ImageScreen
+from tiny_computer_agent.perception import (
+    DetectorUnavailable, IndexedScreen, OcrDetector, detector_available, make_detector,
+)
+
+BOX = [[10, 10], [80, 10], [80, 30], [10, 30]]
+
+
+class FakeDetector:
+    name = "fake"
+
+    def __init__(self, elements=(), error=None):
+        self.elements = tuple(elements)
+        self.error = error
+
+    def detect(self, image):
+        if self.error is not None:
+            raise self.error
+        return self.elements
+
+
+def test_low_confidence_boxes_are_dropped():
+    assert OcrDetector(min_score=0.9)._to_element(BOX, "Save", 0.5) is None
+
+
+def test_prose_fragments_and_punctuation_are_not_widgets():
+    detector = OcrDetector(max_label_chars=20, min_label_chars=2)
+    assert detector._to_element(BOX, "a long sentence that is not a widget", 1.0) is None
+    assert detector._to_element(BOX, "S", 1.0) is None
+    assert detector._to_element(BOX, " ... ", 1.0) is None
+    assert detector._to_element(BOX, " first ", 1.0).label == "first"
+
+
+def test_a_box_becomes_a_clickable_element():
+    element = OcrDetector()._to_element(BOX, "Save", 0.99)
+    assert (element.x, element.y, element.width, element.height) == (10, 10, 70, 20)
+    assert element.center == (45, 20)
+
+
+def test_indexed_screen_replaces_elements_with_detected_ones():
+    screen = IndexedScreen(
+        FakeScreen([(Element("stale"),)]), FakeDetector((Element("Save"), Element("Cancel")))
+    )
+    observation = screen.capture()
+    assert [e.label for e in observation.elements] == ["Save", "Cancel"]
+    assert observation.text == "Save Cancel"          # text filled from the index
+
+
+def test_indexed_screen_keeps_text_from_the_source():
+    screen = IndexedScreen(FakeScreen([(Element("x"),)], text="hello"),
+                           FakeDetector((Element("Save"),)))
+    assert screen.capture().text == "hello"
+
+
+def test_a_perception_error_does_not_kill_the_run():
+    screen = IndexedScreen(FakeScreen([(Element("x"),)]), FakeDetector(error=ValueError("boom")))
+    observation = screen.capture()
+    assert observation.elements == ()
+    assert "perception error" in observation.text
+
+
+def test_a_missing_engine_is_reported_loudly():
+    screen = IndexedScreen(FakeScreen([(Element("x"),)]),
+                           FakeDetector(error=DetectorUnavailable("no engine")))
+    with pytest.raises(DetectorUnavailable):
+        screen.capture()
+
+
+def test_unknown_detector_names_are_reported():
+    available, detail = detector_available("nope")
+    assert available is False and "unknown detector" in detail
+    with pytest.raises(DetectorUnavailable):
+        make_detector("nope")
+
+
+def test_image_screen_replays_a_still_and_holds_it(tmp_path):
+    path = write_fake_desktop(tmp_path / "screen.png")
+    screen = ImageScreen(str(path))
+    assert screen.capture().frame.size == screen.capture().frame.size
+    assert isinstance(screen.capture().text, str)
+
+
+def test_image_screen_rejects_something_empty():
+    with pytest.raises(ValueError):
+        ImageScreen([])
+
+
+def test_the_synthetic_desktop_looks_like_the_intended_screen():
+    image = fake_desktop()
+    assert image.size == (720, 460)
+
+
+@pytest.mark.skipif(not detector_available()[0], reason="no OCR engine installed")
+def test_real_ocr_reads_a_real_screenshot():
+    """The whole point: pixels in, labelled elements with usable coordinates out."""
+    elements = OcrDetector().detect(fake_desktop())
+    labels = [e.label for e in elements]
+    assert len(elements) >= 3, labels
+    assert any(e.matches("Save") for e in elements), labels
+
+    save = next(e for e in elements if e.matches("Save"))
+    # the Save button is drawn at (200..380, 265..300)
+    assert 200 <= save.center[0] <= 380
+    assert 260 <= save.center[1] <= 310
+
+
+@pytest.mark.skipif(not detector_available()[0], reason="no OCR engine installed")
+def test_ocr_index_plugs_straight_into_a_click():
+    """Label -> pixel resolution on coordinates the detector actually produced."""
+    from tiny_computer_agent import Action, resolve_point
+
+    observation = Observation(elements=OcrDetector().detect(fake_desktop()))
+    point = resolve_point(Action("click", label="Save"), observation)
+    assert point is not None
+    assert 200 <= point[0] <= 380