tests/test_share_traces.py

from __future__ import annotations

import importlib.util
import json
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[1]
for candidate in (ROOT / "src", ROOT):
    if str(candidate) not in sys.path:
        sys.path.insert(0, str(candidate))

from app_kit.tracing import write_trace_artifact


def _load_share_module():
    script_path = ROOT / "scripts" / "share_traces_to_hf_dataset.py"
    spec = importlib.util.spec_from_file_location("share_traces_to_hf_dataset", script_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot import {script_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_share_traces_materializes_schema_and_local_artifacts(tmp_path: Path) -> None:
    repo_root = tmp_path / "repo"
    artifact_root = repo_root / "data" / "artifacts" / "p5_memory_quilt"
    trace_path = write_trace_artifact(
        artifact_root,
        {
            "kind": "smoke",
            "project": "p5",
            "pack_id": "p5_memory_quilt",
            "pack_path": "data/demo_packs/p5_memory_quilt",
            "inputs": {"memory_count": 7},
            "parsed_outputs": {
                "tile_count": 7,
                "quilt_path": "artifacts/verification/2026-06-10/p5_smoke/quilt.png",
                "tile_path": "artifacts/verification/2026-06-10/p5_smoke/tile.png",
            },
            "model_name": "jetbrains/Mellm-2-12B-Instruct",
            "model_id": "jetbrains/Mellm-2-12B-Instruct",
            "adapter_name": "local-transformers",
            "checkpoint_path": "/tmp/checkpoint",
            "checkpoint_source": "primary-checkpoint",
            "generation_stats": {
                "prompt_tokens": 11,
                "generated_tokens": 22,
                "elapsed_ms": 12.5,
            },
        },
    )
    assert trace_path.exists()

    module = _load_share_module()
    output_dir = module.default_output_dir(
        repo_root=repo_root,
        repo_name=ROOT.name,
        today="2026-06-10",
    )
    materialized = module.materialize_dataset(
        traces_dir=artifact_root / "traces",
        output_dir=output_dir,
        repo_root=repo_root,
        repo_name=ROOT.name,
        today="2026-06-10",
    )

    assert materialized["records_written"] == 1
    assert materialized["output_dir"] == output_dir
    assert materialized["dataset_jsonl"].exists()
    assert materialized["metadata_json"].exists()

    row = json.loads(materialized["dataset_jsonl"].read_text(encoding="utf-8").strip())
    assert row["repo_name"] == ROOT.name
    assert row["kind"] == "smoke"
    assert row["project"] == "p5"
    assert row["pack_id"] == "p5_memory_quilt"
    assert row["model_name"] == "jetbrains/Mellm-2-12B-Instruct"
    assert row["model_id"] == "jetbrains/Mellm-2-12B-Instruct"
    assert row["adapter_name"] == "local-transformers"
    assert row["checkpoint_source"] == "primary-checkpoint"
    assert row["generation_stats"]["generated_tokens"] == 22
    assert row["parsed_outputs"]["tile_count"] == 7

    metadata = json.loads(materialized["metadata_json"].read_text(encoding="utf-8"))
    assert metadata["repo_name"] == ROOT.name
    assert metadata["schema_fields"] == [
        "timestamp",
        "inputs",
        "parsed_outputs",
        "model_name",
        "model_id",
        "adapter_name",
        "generation_stats",
        "checkpoint_path",
        "checkpoint_source",
    ]
    assert metadata["source_trace_count"] == 1
    assert str(materialized["output_dir"]).endswith(f"artifacts/verification/2026-06-10/sharing_is_caring/{ROOT.name}")