tests/test_trace_stream.py

from __future__ import annotations

import asyncio
from pathlib import Path
from typing import Any

from autocad_bench.orchestration.trace_stream import IncrementalTracePublisher


class FakeS3:
    def __init__(self) -> None:
        self.objects: dict[tuple[str, str], bytes] = {}
        self.puts: list[str] = []

    def put_object(self, **kwargs: Any) -> None:
        key = str(kwargs["Key"])
        self.puts.append(key)
        self.objects[(str(kwargs["Bucket"]), key)] = bytes(kwargs["Body"])

    def upload_file(
        self,
        filename: str,
        bucket: str,
        key: str,
        **_kwargs: Any,
    ) -> None:
        self.puts.append(key)
        self.objects[(bucket, key)] = Path(filename).read_bytes()


def test_incremental_trace_publisher_streams_only_changed_live_artifacts(
    tmp_path: Path,
) -> None:
    root = tmp_path / "run"
    frame = root / "rollouts" / "task-001" / "frames" / "frame_0000.png"
    state = root / "rollouts" / "task-001" / "run-state.json"
    raw_turns = root / "rollouts" / "task-001" / "model_turns.jsonl"
    frame.parent.mkdir(parents=True)
    frame.write_bytes(b"png")
    state.write_text('{"state":"running"}\n', encoding="utf-8")
    raw_turns.write_text('{"turn":1}\n', encoding="utf-8")
    s3 = FakeS3()
    publisher = IncrementalTracePublisher(
        root=root,
        bucket="bucket",
        prefix="controller/runs/benchmark",
        s3=s3,
        interval_s=0.01,
    )

    asyncio.run(publisher.publish_once())

    assert (
        s3.objects[
            (
                "bucket",
                "controller/runs/benchmark/rollouts/task-001/frames/frame_0000.png",
            )
        ]
        == b"png"
    )
    assert not any(key.endswith("model_turns.jsonl") for key in s3.puts)
    first_put_count = len(s3.puts)

    asyncio.run(publisher.publish_once())
    assert len(s3.puts) == first_put_count + 1  # publisher status advanced

    state.write_text('{"state":"completed"}\n', encoding="utf-8")
    asyncio.run(publisher.publish_once(final=True))

    assert any(key.endswith("model_turns.jsonl") for key in s3.puts)
    assert (
        s3.objects[
            (
                "bucket",
                "controller/runs/benchmark/rollouts/task-001/run-state.json",
            )
        ]
        == b'{"state":"completed"}\n'
    )