tests/lightweight/test_TaskGoal.py

"""
test_TaskGoal.py

Coverage tests for all env_id branches in task_goal.get_language_goal.
Run with: uv run python -m pytest tests/lightweight/test_TaskGoal.py -s
"""
from pathlib import Path
import importlib.util
import types

from tests._shared.repo_paths import find_repo_root


def _load_task_goal_module():
    repo_root = find_repo_root(__file__)
    module_path = repo_root / "src" / "robomme" / "robomme_env" / "utils" / "task_goal.py"
    spec = importlib.util.spec_from_file_location("task_goal_under_test", module_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


mod = _load_task_goal_module()
get_language_goal = mod.get_language_goal


class _Unwrapped:
    """Mock env.unwrapped; arbitrary attributes can be set as needed."""
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)


class _Env:
    def __init__(self, unwrapped):
        self.unwrapped = unwrapped


class _Self:
    """Mock self object for called methods."""
    def __init__(self, env, **kwargs):
        self.env = env
        for k, v in kwargs.items():
            setattr(self, k, v)


def _make_self(unwrapped_attrs=None, **self_attrs):
    unwrapped = _Unwrapped(**(unwrapped_attrs or {}))
    env = _Env(unwrapped)
    return _Self(env, **self_attrs)


def _call(env_id, mock_self):
    result = get_language_goal(mock_self, env_id)
    print()
    for idx, goal in enumerate(result, start=1):
        print(f"[{env_id}] goal{idx}: {goal}")
    return result


def test_unknown_env_returns_single_goal_when_equal():
    """Unknown env_id goes to the default empty-string branch; when two goals are equal it returns a single-element list."""
    s = _make_self()
    result = _call("UnknownEnv", s)
    assert result == [""]


def test_movecube_still_returns_two_goals():
    """For defined branches with different texts, still return a two-element list."""
    s = _make_self()
    result = _call("MoveCube", s)
    assert len(result) == 2



def test_binfill_one_color():
    """Branch behavior test case."""
    s = _make_self(unwrapped_attrs=dict(
        red_cubes_target_number=3,
        blue_cubes_target_number=0,
        green_cubes_target_number=0,
    ))
    result = _call("BinFill", s)
    assert "three red cubes" in result[0]
    assert " and " not in result[0]
    assert "three red cubes" in result[1]


def test_binfill_two_colors():
    """Branch behavior test case."""
    s = _make_self(unwrapped_attrs=dict(
        red_cubes_target_number=1,
        blue_cubes_target_number=2,
        green_cubes_target_number=0,
    ))
    result = _call("BinFill", s)
    assert "one red cube" in result[0]
    assert "two blue cubes" in result[0]
    assert " and " in result[0]
    assert "one red cube" in result[1]
    assert "two blue cubes" in result[1]


def test_binfill_three_colors():
    """Branch behavior test case."""
    s = _make_self(unwrapped_attrs=dict(
        red_cubes_target_number=2,
        blue_cubes_target_number=3,
        green_cubes_target_number=1,
    ))
    result = _call("BinFill", s)
    assert "two red cubes" in result[0]
    assert "three blue cubes" in result[0]
    assert "one green cube" in result[0]
    assert "two red cubes" in result[1]
    assert "three blue cubes" in result[1]
    assert "one green cube" in result[1]


# ── PickXtimes: 2 branches ──

def test_pickxtimes_once():
    """Branch behavior test case."""
    s = _make_self(unwrapped_attrs=dict(num_repeats=1, target_color_name="red"))
    result = _call("PickXtimes", s)
    assert "repeating" not in result[0]
    assert "red cube" in result[0]
    if len(result) > 1:
        assert "repeating" not in result[1]
        assert "red cube" in result[1]


def test_pickxtimes_multiple():
    """Branch behavior test case."""
    s = _make_self(unwrapped_attrs=dict(num_repeats=3, target_color_name="blue"))
    result = _call("PickXtimes", s)
    assert "repeating this action three times" in result[0]
    assert "blue cube" in result[0]
    assert "pick-and-place action three times" in result[1]
    assert "blue cube" in result[1]


# ── SwingXtimes: 2 branches ──

def test_swingxtimes_once():
    """Branch behavior test case."""
    s = _make_self(unwrapped_attrs=dict(num_repeats=1, target_color_name="green"))
    result = _call("SwingXtimes", s)
    assert "put it down on the left-side target" in result[0]
    assert "repeating" not in result[0]
    assert "left-side target" in result[1]
    assert "repeating" not in result[1]


def test_swingxtimes_multiple():
    """Branch behavior test case."""
    s = _make_self(unwrapped_attrs=dict(num_repeats=5, target_color_name="red"))
    result = _call("SwingXtimes", s)
    assert "repeating this back and forth motion five times" in result[0]
    assert "right-to-left swing motion five times" in result[1]


# ── VideoUnmask: 2 branches ──

def test_videounmask_pick_one():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(
            color_names=["red", "blue", "green"],
            configs={"easy": {"pick": 1}},
        ),
        difficulty="easy",
    )
    result = _call("VideoUnmask", s)
    assert "red cube" in result[0]
    assert "another container" not in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "red cube" in g
    assert "another container" not in g


def test_videounmask_pick_two():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(
            color_names=["red", "blue", "green"],
            configs={"hard": {"pick": 2}},
        ),
        difficulty="hard",
    )
    result = _call("VideoUnmask", s)
    assert "red cube" in result[0]
    assert "another container hiding the blue cube" in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "red cube" in g
    assert "another container hiding the blue cube" in g


# ── VideoUnmaskSwap: 2 branches ──

def test_videounmaskswap_pick_one():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(color_names=["red", "blue", "green"]),
        pick_times=1,
    )
    result = _call("VideoUnmaskSwap", s)
    assert "red cube" in result[0]
    assert "another container" not in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "red cube" in g
    assert "another container" not in g


def test_videounmaskswap_pick_two():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(color_names=["red", "blue", "green"]),
        pick_times=2,
    )
    result = _call("VideoUnmaskSwap", s)
    assert "red cube" in result[0]
    assert "another container hiding the blue cube" in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "red cube" in g
    assert "another container hiding the blue cube" in g


# ── ButtonUnmask: 2 branches ──

def test_buttonunmask_pick_one():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(
            color_names=["green", "red", "blue"],
            configs={"easy": {"pick": 1}},
        ),
        difficulty="easy",
    )
    result = _call("ButtonUnmask", s)
    assert "press the button" in result[0]
    assert "green cube" in result[0]
    assert "another container" not in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "press the button" in g
    assert "green cube" in g
    assert "another container" not in g


def test_buttonunmask_pick_two():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(
            color_names=["green", "red", "blue"],
            configs={"hard": {"pick": 2}},
        ),
        difficulty="hard",
    )
    result = _call("ButtonUnmask", s)
    assert "press the button" in result[0]
    assert "another container hiding the red cube" in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "press the button" in g
    assert "another container hiding the red cube" in g


# ── ButtonUnmaskSwap: 2 branches ──

def test_buttonunmaskswap_pick_one():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(color_names=["blue", "green", "red"]),
        pick_times=1,
    )
    result = _call("ButtonUnmaskSwap", s)
    assert "press both buttons" in result[0]
    assert "blue cube" in result[0]
    assert "another container" not in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "press both buttons" in g
    assert "blue cube" in g
    assert "another container" not in g


def test_buttonunmaskswap_pick_two():
    """Branch behavior test case."""
    s = _make_self(
        unwrapped_attrs=dict(color_names=["blue", "green", "red"]),
        pick_times=2,
    )
    result = _call("ButtonUnmaskSwap", s)
    assert "press both buttons" in result[0]
    assert "another container hiding the green cube" in result[0]
    g = result[1] if len(result) > 1 else result[0]
    assert "press both buttons" in g
    assert "another container hiding the green cube" in g


# ── VideoPlaceButton: 1 branch ──

def test_videoplacebutton():
    """Branch behavior test case."""
    s = _make_self(
        target_color_name="red",
        target_target_language="after",
    )
    result = _call("VideoPlaceButton", s)
    assert "red cube" in result[0]
    assert "right after the button was pressed" in result[0]
    assert "red cube" in result[1]
    assert "where it was placed immediately after the button was pressed" in result[1]


# ── VideoPlaceOrder: 1 branch ──

def test_videoplaceorder():
    """Case: place at the N-th target in order.Inputs: target_color_name=blue, which_in_subset=3."""
    s = _make_self(
        target_color_name="blue",
        which_in_subset=3,
    )
    result = _call("VideoPlaceOrder", s)
    assert "blue cube" in result[0]
    assert "third target" in result[0]
    assert "blue cube" in result[1]
    assert "third target" in result[1]
    assert "where it was placed" in result[1]


# ── PickHighlight: 1 branch ──

def test_pickhighlight():
    """Branch behavior test case."""
    s = _make_self()
    result = _call("PickHighlight", s)
    assert "press the button" in result[0]
    assert "highlighteted" in result[0]
    assert "highlighted cubes" in result[1]
    assert "press the button again to stop" in result[1]


# ── VideoRepick: 2 branches ──
# Inputs: self.num_repeats.

def test_videorepick_once():
    """Branch behavior test case."""
    s = _make_self(num_repeats=1)
    result = _call("VideoRepick", s)
    assert "pick up the same block" in result[0]
    assert "repeatedly" not in result[0]
    assert "pick up the same cube" in result[1]
    assert "repeatedly" not in result[1]


def test_videorepick_multiple():
    """Branch behavior test case."""
    s = _make_self(num_repeats=4)
    result = _call("VideoRepick", s)
    assert "repeatedly pick up and put down" in result[0]
    assert "four times" in result[0]
    assert "same cube" in result[1]
    assert "four times" in result[1]


# ── StopCube: 1 branch ──

def test_stopcube():
    """Case: press button to stop at the N-th arrival.Inputs: stop_time=2."""
    s = _make_self(unwrapped_attrs=dict(stop_time=2))
    result = _call("StopCube", s)
    assert "second time" in result[0]
    assert "second visit" in result[1]


# ── InsertPeg: 1 branch ──

def test_insertpeg():
    """Case: insert the same peg end into the same side.Inputs: none."""
    s = _make_self()
    result = _call("InsertPeg", s)
    assert "grasp the same end" in result[0]
    assert "grasp the same peg at the same end" in result[1]
    assert "as in the video" in result[1]


# ── MoveCube: 1 branch ──

def test_movecube():
    """Case: move the cube to target in the same way as before.Inputs: none."""
    s = _make_self()
    result = _call("MoveCube", s)
    assert "move the cube to the target" in result[0]
    assert "shown in the video" in result[1]


# ── PatternLock: 1 branch ──

def test_patternlock():
    """Case: retrace the same pattern with the stick.Inputs: none."""
    s = _make_self()
    result = _call("PatternLock", s)
    assert "retrace the same pattern" in result[0]
    assert "retrace the same pattern shown in the video" in result[1]


# ── RouteStick: 1 branch ──

def test_routestick():
    """Case: use the stick to follow the same path around the rods.Inputs: none."""
    s = _make_self()
    result = _call("RouteStick", s)
    assert "navigate around the sticks" in result[0]
    assert "following the same path shown in the video" in result[1]