tests/test_models.py

"""TDD tests for kinchat.models — Pydantic schemas."""
import json
import pytest
from pydantic import ValidationError

from kinchat.models import (
    ChatMsg,
    KinChatAction,
    KinChatObservation,
    KinChatState,
    PersonaState,
    Secret,
)


# ---- KinChatAction ----

def test_action_valid_send():
    a = KinChatAction(action_type="send", message="hi", recipients=["mom"])
    assert a.action_type == "send"
    assert a.recipients == ["mom"]


def test_action_stay_silent_defaults():
    a = KinChatAction(action_type="stay_silent")
    assert a.message == ""
    assert a.recipients == []
    assert a.reasoning == ""


def test_action_rejects_invalid_type():
    with pytest.raises(ValidationError):
        KinChatAction(action_type="shout")


@pytest.mark.parametrize("t", ["send", "edit", "block", "suggest", "stay_silent"])
def test_action_all_five_types(t):
    KinChatAction(action_type=t)


def test_action_json_roundtrip():
    a = KinChatAction(
        action_type="edit",
        message="softer phrasing",
        recipients=["grandma"],
        reasoning="avoid worry",
    )
    blob = a.model_dump_json()
    b = KinChatAction.model_validate_json(blob)
    assert b == a


# ---- ChatMsg ----

def test_chatmsg_basic():
    m = ChatMsg(sender="mom", text="dinner?", recipients=["user"], turn=3)
    assert m.turn == 3
    assert m.recipients == ["user"]


def test_chatmsg_group_default_recipients():
    m = ChatMsg(sender="sib1", text="hi all", turn=0)
    assert m.recipients == []  # empty = group chat


# ---- PersonaState ----

def test_persona_state_defaults():
    p = PersonaState(persona_id="mom")
    assert 0.0 <= p.trust <= 1.0
    assert p.mood in {"calm", "annoyed", "worried", "warm"}
    assert isinstance(p.known_facts, set)


def test_persona_state_known_facts_serializes_as_list():
    p = PersonaState(persona_id="dad", known_facts={"fact_1", "fact_2"})
    blob = json.loads(p.model_dump_json())
    assert isinstance(blob["known_facts"], list)
    assert set(blob["known_facts"]) == {"fact_1", "fact_2"}


def test_persona_state_roundtrip_set_preserved():
    p = PersonaState(persona_id="sib1", known_facts={"a", "b"})
    p2 = PersonaState.model_validate_json(p.model_dump_json())
    assert p2.known_facts == {"a", "b"}


def test_persona_invalid_mood_rejected():
    with pytest.raises(ValidationError):
        PersonaState(persona_id="x", mood="enraged")


def test_persona_trust_range_enforced():
    with pytest.raises(ValidationError):
        PersonaState(persona_id="x", trust=1.5)
    with pytest.raises(ValidationError):
        PersonaState(persona_id="x", trust=-0.1)


# ---- Secret ----

def test_secret_basic():
    s = Secret(
        fact_id="grandma_surgery",
        text="grandma is having knee surgery",
        owners=["mom", "dad"],
        sensitivity=0.8,
    )
    assert s.sensitivity == 0.8


def test_secret_sensitivity_range():
    with pytest.raises(ValidationError):
        Secret(fact_id="x", text="y", owners=[], sensitivity=1.2)


# ---- KinChatState ----

def test_kinchat_state_construction():
    st = KinChatState(
        session_id="s1",
        episode_index=0,
        scenario_id="secret_grandma_01",
        family={"mom": PersonaState(persona_id="mom")},
        secrets=[Secret(fact_id="f", text="t", owners=["mom"], sensitivity=0.5)],
        chat_log=[],
        turn_index=0,
        scores_history=[],
    )
    assert st.episode_index == 0
    assert "mom" in st.family


def test_kinchat_state_roundtrip():
    st = KinChatState(
        session_id="s1",
        episode_index=2,
        scenario_id="sc",
        family={"mom": PersonaState(persona_id="mom", known_facts={"f1"})},
        secrets=[],
        chat_log=[ChatMsg(sender="mom", text="hi", turn=0)],
        turn_index=1,
        scores_history=[{"leak": 0.0, "audience_fit": 0.7}],
    )
    st2 = KinChatState.model_validate_json(st.model_dump_json())
    assert st2.family["mom"].known_facts == {"f1"}
    assert st2.chat_log[0].text == "hi"


# ---- KinChatObservation ----

def test_observation_construction():
    obs = KinChatObservation(
        chat_history=[ChatMsg(sender="mom", text="hi", turn=0)],
        user_draft="hey mom",
        active_recipients=["mom"],
        turn_index=0,
        scenario_brief="mom asks how work is going",
        reward=0.42,
        reward_breakdown={"leak": 0.0, "audience_fit": 0.8, "restraint": 0.5, "trust_delta": 0.1},
        feedback="ok",
        done=False,
        episode_id="ep_0",
    )
    assert obs.done is False
    assert obs.reward == 0.42
    assert obs.reward_breakdown["audience_fit"] == 0.8


def test_observation_roundtrip():
    obs = KinChatObservation(
        chat_history=[],
        user_draft="",
        active_recipients=[],
        turn_index=0,
        scenario_brief="",
        reward=0.0,
        reward_breakdown={},
        feedback="",
        done=True,
        episode_id="ep_X",
    )
    obs2 = KinChatObservation.model_validate_json(obs.model_dump_json())
    assert obs2 == obs