tests/test_publish_trace.py

"""One-click 'Publish my trace to the Hub' (Sharing is Caring). Network-free."""
import os

os.environ.setdefault("USE_STUB_EXTRACTOR", "1")

import huggingface_hub  # noqa: E402

from server import events as bus  # noqa: E402
from ui import blocks  # noqa: E402


def test_publish_without_token_is_graceful_and_offline(monkeypatch):
    monkeypatch.delenv("HF_WRITE_TOKEN", raising=False)

    def _boom(*a, **k):  # must NOT touch the network without a token
        raise AssertionError("HfApi must not be constructed without a token")

    monkeypatch.setattr(huggingface_hub, "HfApi", _boom)
    msg = blocks._on_publish_trace()
    assert "HF_WRITE_TOKEN" in msg and "isn't configured" in msg


def test_publish_empty_run(monkeypatch):
    monkeypatch.setenv("HF_WRITE_TOKEN", "hf_dummy")
    bus.reset() if hasattr(bus, "reset") else None
    # no run recorded → export_run has no steps
    import server.trace as trace
    monkeypatch.setattr(trace, "export_run", lambda redact=True: {"steps": [], "summary": {}})
    msg = blocks._on_publish_trace()
    assert "Nothing to publish" in msg


def test_publish_with_token_uploads_public(monkeypatch):
    monkeypatch.setenv("HF_WRITE_TOKEN", "hf_dummy")
    monkeypatch.setenv("TRACE_DATASET", "me/traces")
    import server.trace as trace
    monkeypatch.setattr(trace, "export_run",
                        lambda redact=True: {"steps": [{"stage": "model"}], "summary": {"steps": 1}})

    calls = {}

    class FakeApi:
        def __init__(self, token=None):
            calls["token"] = token

        def create_repo(self, repo_id, repo_type=None, private=None, exist_ok=None):
            calls["create"] = {"repo_id": repo_id, "repo_type": repo_type, "private": private}

        def upload_file(self, **k):
            calls["upload"] = k

    monkeypatch.setattr(huggingface_hub, "HfApi", FakeApi)
    msg = blocks._on_publish_trace()
    assert calls["token"] == "hf_dummy"
    assert calls["create"] == {"repo_id": "me/traces", "repo_type": "dataset", "private": False}
    assert calls["upload"]["repo_id"] == "me/traces" and calls["upload"]["repo_type"] == "dataset"
    assert "me/traces" in msg and "huggingface.co/datasets/me/traces" in msg