import pandas as pd

from arena.ui.session_cache import (
    add_run_to_cache,
    comparison_table,
    cached_loss_curves,
)


def test_session_cache_stores_summary_and_history():
    history = pd.DataFrame([{"step": 0, "loss": 2.0}, {"step": 1, "loss": 1.0}])
    summary = {
        "task": "Linear Regression",
        "optimizer": "ToyOpt",
        "device": "cpu",
        "seed": 0,
        "steps": 2,
        "initial_loss": 2.0,
        "final_loss": 1.0,
        "best_loss": 1.0,
        "time_sec": 0.1,
        "code_hash": "abc",
        "status": "success",
    }
    cache = add_run_to_cache([], summary=summary, history=history, pseudocode="pseudo")
    table = comparison_table(cache)
    curves = cached_loss_curves(cache)
    assert len(cache) == 1
    assert table.loc[0, "best_loss"] == 1.0
    assert list(curves["loss"]) == [2.0, 1.0]