tests/test_vlm_task_generator.py

from __future__ import annotations

from pathlib import Path

from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.tasks.validators import validate_task
from dovla_cil.utils.io import iter_jsonl
from dovla_cil.vlm.client import VLMClient
from dovla_cil.vlm.task_generator import (
    TaskGenerator,
    default_task_generation_request,
)


class FakeTaskClient(VLMClient):
    def __init__(self, responses):
        super().__init__(api_key="fake", model="fake")
        self.responses = list(responses)
        self.calls = 0

    def chat_json(self, system, user, schema_hint=None):
        del system, user, schema_hint
        self.calls += 1
        if not self.responses:
            raise AssertionError("Unexpected VLM call")
        return self.responses.pop(0)


def test_mock_vlm_returns_valid_tasks(monkeypatch) -> None:
    monkeypatch.setenv("OPENCLAUDE_MOCK", "1")
    request = default_task_generation_request(num_tasks=3)
    tasks = TaskGenerator(VLMClient(api_key="fake", model="mock")).generate_tasks(request)

    assert len(tasks) == 3
    for task in tasks:
        validate_task(task)


def test_invalid_generated_task_is_repaired() -> None:
    valid = built_in_toy_tasks()[0].to_dict()
    invalid = dict(valid)
    invalid["success_predicates"] = [{"name": "grasped", "args": ["missing"]}]
    client = FakeTaskClient([{"tasks": [invalid]}, valid])
    request = default_task_generation_request(num_tasks=1)

    tasks = TaskGenerator(client).generate_tasks(request)

    assert len(tasks) == 1
    assert tasks[0].task_id == valid["task_id"]
    assert client.calls == 2


def test_invalid_generated_task_is_skipped_when_repair_fails() -> None:
    valid = built_in_toy_tasks()[0].to_dict()
    invalid = dict(valid)
    invalid["success_predicates"] = [{"name": "grasped", "args": ["missing"]}]
    still_invalid = dict(invalid)
    client = FakeTaskClient([{"tasks": [invalid]}, still_invalid])

    tasks = TaskGenerator(client).generate_tasks(default_task_generation_request(num_tasks=1))

    assert tasks == []
    assert client.calls == 2


def test_script_runs_in_mock_mode_without_network(tmp_path: Path, monkeypatch) -> None:
    from scripts.generate_tasks import main

    out = tmp_path / "tasks.jsonl"
    monkeypatch.delenv("OPENCLAUDE_API_KEY", raising=False)
    monkeypatch.delenv("OPENCLAUDE_MODEL", raising=False)

    result = main(["--num-tasks", "2", "--out", str(out), "--mock", "--seed", "5"])

    assert result == 0
    rows = list(iter_jsonl(out))
    assert len(rows) == 2
    assert rows[0]["task_id"]