tests/test_task_prompt.py

"""Unit tests for fpgm.prompting.task_prompt -- pure stdlib, no GPU or weights."""

from __future__ import annotations

import pytest

from fpgm.config import SegmentationConfig
from fpgm.prompting.task_prompt import PromptCandidates, TaskPromptDeriver

# (instruction, expected primary, expected alternatives)
CASES = [
    (
        "Put brick in drawer shelf and close drawer",
        "brick",
        ["drawer shelf", "drawer"],
    ),
    (
        # "pick up" particle: "up" must not be swallowed into the object phrase.
        "Pick up the marker and put it in the cup",
        "marker",
        ["cup"],
    ),
    (
        # leading article ("the") right after the verb.
        "Open the top drawer",
        "top drawer",
        [],
    ),
    (
        "Push the red block onto the table",
        "red block",
        ["table"],
    ),
    (
        "Take the bowl from the shelf and place it on the table",
        "bowl",
        ["shelf", "table"],
    ),
    (
        # two-word connector "next to".
        "Move the cup next to the plate",
        "cup",
        ["plate"],
    ),
    (
        "Close the drawer",
        "drawer",
        [],
    ),
    (
        "Insert the plug into the socket",
        "plug",
        ["socket"],
    ),
]


@pytest.mark.parametrize("instruction,expected_primary,expected_alternatives", CASES)
def test_derive_primary_and_alternatives(
    instruction: str, expected_primary: str, expected_alternatives: list[str]
) -> None:
    result = TaskPromptDeriver().derive(instruction)
    assert result.primary == expected_primary
    assert result.alternatives == expected_alternatives
    assert result.source == "heuristic"


def test_override_wins_outright() -> None:
    deriver = TaskPromptDeriver()
    result = deriver.derive("Put brick in drawer shelf and close drawer", override="gripper target")
    assert result.primary == "gripper target"
    assert result.alternatives == []
    assert result.source == "override"


def test_config_level_override_by_episode_uuid() -> None:
    config = SegmentationConfig(prompt_overrides={"ep-1": "widget"})
    deriver = TaskPromptDeriver(config=config)

    overridden = deriver.derive("Put brick in drawer shelf and close drawer", episode_uuid="ep-1")
    assert overridden.primary == "widget"
    assert overridden.source == "override"

    # An episode uuid with no matching override falls back to the heuristic.
    not_overridden = deriver.derive("Close the drawer", episode_uuid="some-other-episode")
    assert not_overridden.primary == "drawer"
    assert not_overridden.source == "heuristic"


def test_explicit_override_beats_config_override() -> None:
    config = SegmentationConfig(prompt_overrides={"ep-1": "widget"})
    deriver = TaskPromptDeriver(config=config)
    result = deriver.derive("Close the drawer", override="explicit wins", episode_uuid="ep-1")
    assert result.primary == "explicit wins"


def test_prompt_candidates_all_property_orders_primary_first() -> None:
    candidates = PromptCandidates(primary="brick", alternatives=["drawer shelf", "drawer"])
    assert candidates.all == ["brick", "drawer shelf", "drawer"]


def test_empty_instruction_raises() -> None:
    with pytest.raises(ValueError):
        TaskPromptDeriver().derive("   ")


def test_instruction_without_leading_verb_still_extracts_object() -> None:
    # No recognised leading verb: the heuristic should not crash, and should
    # still split on connectors starting from the first token.
    result = TaskPromptDeriver().derive("Marker into cup")
    assert result.primary == "marker"
    assert result.alternatives == ["cup"]