tests/test_inference_logging.py

"""
Test that inference.py produces the exact stdout format required.
Runs the environment directly (no LLM needed) with mock actions to verify logging format.
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from triage_flow.environment import TriageEnvironment
from triage_flow.graders import grade_task
from models import TriageAction, ActionType, PriorityLevel

# Test the logging functions from inference.py
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from inference import log_start, log_step, log_end

import io
from contextlib import redirect_stdout

# Capture stdout
f = io.StringIO()
with redirect_stdout(f):
    log_start(task="basic-triage", env="triage-flow", model="test-model")

    # Simulate a 3-step episode
    log_step(step=1, action="assign_priority(P001,immediate)", reward=0.50, done=False, error=None)
    log_step(step=2, action="assign_priority(P002,non_urgent)", reward=0.10, done=False, error=None)
    log_step(step=3, action="assign_priority(P003,urgent)", reward=0.50, done=True, error=None)

    log_end(success=True, steps=3, score=1.00, rewards=[0.50, 0.10, 0.50])

output = f.getvalue()
lines = output.strip().split("\n")

print("=== CAPTURED STDOUT ===")
for line in lines:
    print(repr(line))
print()

# Validate format
assert lines[0].startswith("[START]"), f"Line 0 should start with [START]: {lines[0]}"
assert "task=basic-triage" in lines[0], f"Missing task: {lines[0]}"
assert "env=triage-flow" in lines[0], f"Missing env: {lines[0]}"
assert "model=test-model" in lines[0], f"Missing model: {lines[0]}"

for i in range(1, 4):
    assert lines[i].startswith("[STEP]"), f"Line {i} should start with [STEP]: {lines[i]}"
    assert f"step={i}" in lines[i], f"Missing step={i}: {lines[i]}"
    assert "reward=" in lines[i], f"Missing reward: {lines[i]}"
    assert "done=" in lines[i], f"Missing done: {lines[i]}"
    assert "error=" in lines[i], f"Missing error: {lines[i]}"

# Check [STEP] format details
assert "done=false" in lines[1], f"done should be lowercase: {lines[1]}"
assert "done=true" in lines[3], f"done should be lowercase: {lines[3]}"
assert "error=null" in lines[1], f"error should be 'null': {lines[1]}"
assert "reward=0.50" in lines[1], f"reward should be 2 decimal: {lines[1]}"

# Check [END]
assert lines[4].startswith("[END]"), f"Line 4 should start with [END]: {lines[4]}"
assert "success=true" in lines[4], f"success should be lowercase: {lines[4]}"
assert "steps=3" in lines[4], f"Missing steps: {lines[4]}"
assert "score=1.00" in lines[4], f"score should be 2 decimal: {lines[4]}"
assert "rewards=0.50,0.10,0.50" in lines[4], f"rewards format wrong: {lines[4]}"

print("=== ALL FORMAT CHECKS PASSED ===")