Files
Video-Tree-TRM5/tests/unit/test_baseline_run_log.py
T

34 lines
1.2 KiB
Python

import json
import pytest
from app.harness.baseline_run_log import StepsJsonRunLog
class _FakeInner:
def __init__(self, preds, traces):
self._preds, self._traces = preds, traces
async def get_predictions(self, run_id, *, question_ids=None):
return [p for p in self._preds if not question_ids or p["question_id"] in question_ids]
async def get_traces(self, run_id, *, question_ids=None):
return list(self._traces)
@pytest.mark.asyncio
async def test_get_traces_falls_back_to_steps_json_when_table_empty():
steps = [{"thought": "t", "tool_call": {"tool": "view_node", "args": {}}, "tool_output": "o"}]
preds = [{"video_id": "v1", "question_id": "q1", "steps_json": json.dumps(steps)}]
log = StepsJsonRunLog(_FakeInner(preds, traces=[]))
rows = await log.get_traces("r", question_ids=["q1"])
assert rows[0]["tool_name"] == "view_node" and rows[0]["question_id"] == "q1"
@pytest.mark.asyncio
async def test_get_traces_prefers_nonempty_inner_table():
inner_traces = [{"video_id": "v1", "question_id": "q1", "step": 0, "tool_name": "x"}]
log = StepsJsonRunLog(_FakeInner([], inner_traces))
rows = await log.get_traces("r")
assert rows == inner_traces