tests/conftest.py

"""Shared fixtures: an offline vector store, a stand-in G1 model, and an app wired to a scripted LLM."""

from __future__ import annotations

import hashlib
import re
from dataclasses import replace
from pathlib import Path

import numpy as np
import pytest
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
from fastapi.testclient import TestClient

from api import keys
from api.auth import Allowlist
from api.config import get_settings
from api.llm import ScriptedLLM
from api.logstore import LogStore
from api.main import Services, create_app
from api.runpod_client import RunPodError
from api.storage import LocalStorage
from rag.chunking import Chunk
from rag.store import RagStore

FAKE_G1 = """<mujoco model="fake g1">
  <compiler angle="radian"/>
  <worldbody>
    <body name="pelvis" pos="0 0 0.8">
      <freejoint name="floating_base_joint"/>
      <geom name="torso" type="capsule" size="0.1 0.3" mass="10"/>
      <body name="left_knee_link" pos="0 0.1 -0.4">
        <joint name="left_knee_joint" axis="0 1 0" range="0 2.8"/>
        <geom type="capsule" size="0.05 0.2"/>
      </body>
    </body>
  </worldbody>
  <actuator><motor name="left_knee_joint" joint="left_knee_joint"/></actuator>
  <keyframe><key name="stand" qpos="0 0 0.8 1 0 0 0 0"/></keyframe>
</mujoco>
"""

GOOD_SCENE = """<mujoco model="g1 on slope">
  <include file="g1.xml"/>
  <worldbody>
    <light name="sun" pos="0 0 4"/>
    <geom name="floor" type="plane" size="0 0 0.05" friction="1 0.005 0.0001"/>
    <geom name="slope" type="box" size="1 1 0.05" pos="2 0 0.2" euler="0 -0.26 0"/>
    <body name="crate" pos="1 1 0.3"><freejoint/><geom name="crate_geom" type="box" size="0.2 0.2 0.2" mass="3"/></body>
  </worldbody>
</mujoco>"""


class HashEmbedding(EmbeddingFunction[Documents]):
    """Bag-of-words hashing embedding: deterministic, offline, good enough to rank by word overlap."""

    def __init__(self) -> None:
        pass

    def __call__(self, input: Documents) -> Embeddings:
        vectors = []
        for text in input:
            v = np.zeros(256, dtype=np.float32)
            for word in re.findall(r"[a-z0-9_]+", text.lower()):
                v[int(hashlib.md5(word.encode()).hexdigest(), 16) % 256] += 1.0
            n = np.linalg.norm(v)
            vectors.append(v / n if n else v)
        return vectors

    @staticmethod
    def name() -> str:
        return "test-hash"

    def get_config(self) -> dict:
        return {}

    @staticmethod
    def build_from_config(config: dict) -> HashEmbedding:
        return HashEmbedding()


@pytest.fixture
def store(tmp_path) -> RagStore:
    s = RagStore(tmp_path / "chroma", embedding_function=HashEmbedding())
    s.add(
        "mujoco",
        [
            Chunk(
                "[XML reference > geom] geom type plane box sphere friction size pos euler", "doc#geom", {"kind": "doc"}
            ),
            Chunk("[XML reference > light] light pos dir diffuse", "doc#light", {"kind": "doc"}),
            Chunk(
                '[MJCF example scene.xml] <include file="g1.xml"/> floor plane',
                "menagerie:scene.xml",
                {"kind": "example", "robot": "unitree_g1"},
            ),
        ],
    )
    s.add(
        "ros2",
        [
            Chunk(
                "[Writing a simple py publisher] create_publisher String timer_callback rclpy",
                "docs:pubsub",
                {"kind": "doc", "distro": "humble"},
            ),
            Chunk(
                "[Launch file formats] LaunchDescription Node remappings parameters",
                "docs:launch",
                {"kind": "doc", "distro": "humble"},
            ),
            Chunk(
                "[ROS 2 humble Message sensor_msgs/msg/LaserScan]\nfloat32[] ranges",
                "ci:LaserScan.msg",
                {"kind": "interface", "distro": "humble", "interface": "sensor_msgs/msg/LaserScan"},
            ),
        ],
    )
    return s


@pytest.fixture
def menagerie(tmp_path) -> Path:
    root = tmp_path / "menagerie"
    (root / "unitree_g1").mkdir(parents=True)
    (root / "unitree_g1" / "g1.xml").write_text(FAKE_G1)
    return root


@pytest.fixture
def settings(tmp_path, menagerie):
    base = get_settings()
    return replace(
        base,
        llm_api_key="test",
        data_dir=tmp_path / "data",
        testers_file=tmp_path / "testers.yaml",
        log_db=tmp_path / "data" / "log.sqlite3",
        menagerie_dir=menagerie,
        mjcf_sim_seconds=0.2,
        codegen_backend="direct",
        auto_review=False,
        prompt_enhance=False,
        free_debug=False,
    )


# --------------------------------------------------------------------------- job backend


class FakeBackend:
    name = "runpod"

    def __init__(self):
        self.jobs: dict[str, dict] = {}
        self.submitted: list[dict] = []
        self.fail_submit = False
        self.functions: list[str] = []

    def supports(self, function):
        return True

    def submit(self, job_input, execution_timeout_s, function="run_simulation"):
        self.functions.append(function)
        if self.fail_submit:
            raise RunPodError("RunPod /run returned HTTP 500: boom")
        job_id = f"job-{len(self.submitted) + 1}"
        self.submitted.append(job_input)
        self.jobs[job_id] = {"id": job_id, "status": "IN_QUEUE"}
        return dict(self.jobs[job_id])

    def status(self, job_id):
        return dict(self.jobs[job_id])


@pytest.fixture
def run_env(settings, store, tmp_path):
    alice = keys.add(settings.testers_file, "alice")
    bob = keys.add(settings.testers_file, "bob")
    llm = ScriptedLLM([])
    backend = FakeBackend()
    storage = LocalStorage(tmp_path / "r2", base_url="https://r2.example/bucket")
    services = Services(
        settings=settings,
        store=store,
        logs=LogStore(settings.log_db),
        allowlist=Allowlist(settings.testers_file),
        llm_factory=lambda: llm,
        backend=backend,
        storage=storage,
    )
    client = TestClient(create_app(services))
    llm.responses[:] = [f"```xml\n{GOOD_SCENE}\n```"]
    gen = client.post(
        "/simulate/generate", json={"description": "G1 near a slope and a crate"}, headers={"X-API-Key": alice}
    ).json()
    return client, {"alice": alice, "bob": bob}, backend, services, gen["id"]


def complete(backend, job_id, run_id, **summary):
    backend.jobs[job_id] = {
        "id": job_id,
        "status": "COMPLETED",
        "executionTime": 12500,
        "delayTime": 3000,
        "output": {
            "run_id": run_id,
            "summary": {"simulated_seconds": 10.0, "fell": False, **summary},
            "files": {
                "mcap": f"runs/{run_id}/run.mcap",
                "video": f"runs/{run_id}/video.mp4",
                "video_webm": f"runs/{run_id}/video.webm",
                "summary": f"runs/{run_id}/summary.json",
            },
        },
    }