CODE_ASSISTANT/5_django_backend/tests/test_stream.py

"""POST /chat/stream — OpenCode 이벤트 → 프론트 SSE 계약(token/title/usage/done/error)."""

import asyncio
import json

import pytest
from django.test import AsyncClient

from apps.chat import events as events_mod
from apps.chat import stream as stream_mod
from apps.chat.models import ChatMessage, ChatSession
from apps.chat.stream import TurnState


# ── TurnState 단위 ───────────────────────────────────────────────
def _part_updated(mid, pid, text, delta=None, ptype="text", sid="s1"):
    props = {"part": {"id": pid, "messageID": mid, "sessionID": sid, "type": ptype, "text": text}}
    if delta is not None:
        props["delta"] = delta
    return {"type": "message.part.updated", "properties": props}


def _msg_updated(mid, role, sid="s1", error=None):
    info = {"id": mid, "sessionID": sid, "role": role}
    if error:
        info["error"] = error
    return {"type": "message.updated", "properties": {"info": info}}


def test_turnstate_delta_path():
    st = TurnState("s1", "q")
    st.handle(_msg_updated("mu", "user"))
    st.handle(_part_updated("mu", "pu", "q"))  # 사용자 echo → 무시
    st.handle(_msg_updated("ma", "assistant"))
    assert st.handle(_part_updated("ma", "pa", "안", delta="안")) == [("token", {"delta": "안"})]
    assert st.handle(_part_updated("ma", "pa", "안녕", delta="녕")) == [("token", {"delta": "녕"})]
    assert st.accumulated() == "안녕"


def test_turnstate_snapshot_path_without_delta():
    st = TurnState("s1", "q")
    st.handle(_msg_updated("ma", "assistant"))
    assert st.handle(_part_updated("ma", "pa", "ab")) == [("token", {"delta": "ab"})]
    assert st.handle(_part_updated("ma", "pa", "abcd")) == [("token", {"delta": "cd"})]
    assert st.handle(_part_updated("ma", "pa", "abcd")) == []  # 같은 스냅샷 재전송
    assert st.accumulated() == "abcd"


def test_turnstate_ignores_reasoning_and_emits_title():
    st = TurnState("s1", "q")
    st.handle(_msg_updated("ma", "assistant"))
    assert st.handle(_part_updated("ma", "pr", "생각중", ptype="reasoning")) == []
    ev = {"type": "session.updated", "properties": {"info": {"id": "s1", "title": "MARA 조회"}}}
    assert st.handle(ev) == [("title", {"title": "MARA 조회"})]
    assert st.handle(ev) == []  # 같은 제목 반복 안 보냄
    placeholder = {"type": "session.updated", "properties": {"info": {"id": "s1", "title": "New session - 2026-09-16T08:54:31.272Z"}}}
    assert st.handle(placeholder) == []  # OpenCode 기본 제목은 무시


def _delta(mid, pid, delta, sid="s1", field="text"):
    return {"type": "message.part.delta", "properties": {"sessionID": sid, "messageID": mid, "partID": pid, "field": field, "delta": delta}}


def test_turnstate_real_server_shape_delta_events():
    """실서버 1.18.6 순서: part.updated(타입, 빈 텍스트) → message.part.delta ×N → part.updated(최종 스냅샷)."""
    st = TurnState("s1", "q")
    st.handle(_msg_updated("mu", "user"))
    st.handle(_part_updated("mu", "pu", "q"))
    st.handle(_msg_updated("ma", "assistant"))
    assert st.handle(_part_updated("ma", "pr", "", ptype="reasoning")) == []
    assert st.handle(_delta("ma", "pr", "생각")) == []  # reasoning delta 는 안 보냄
    assert st.handle(_part_updated("ma", "pt", "")) == []  # 빈 스냅샷
    assert st.handle(_delta("ma", "pt", "SELECT")) == [("token", {"delta": "SELECT"})]
    assert st.handle(_delta("ma", "pt", " SINGLE")) == [("token", {"delta": " SINGLE"})]
    assert st.handle(_part_updated("ma", "pt", "SELECT SINGLE")) == []  # 최종 스냅샷 = 누적 → 중복 없음
    assert st.handle(_part_updated("ma", "pt", "SELECT SINGLE mtart")) == [("token", {"delta": " mtart"})]  # 못 받은 꼬리 보충
    assert st.handle(_delta("ma", "unknown", "x")) == []  # 타입 모르는 파트
    assert st.accumulated() == "SELECT SINGLE mtart"


# ── 통합: 뷰 + 이벤트 버스 + 가짜 OpenCode ───────────────────────
class FakeOpencode:
    def __init__(self, script):
        self.script = script  # prompt 뒤에 흘려보낼 이벤트 dict 목록
        self.prompted = asyncio.Event()
        self.payloads = []
        self.aborted = []
        self.fail_prompt = False
        self.messages = []
        self.session_info = {"id": "s1", "title": ""}

    async def prompt_async(self, sid, payload):
        self.payloads.append((sid, payload))
        if self.fail_prompt:
            raise RuntimeError("down")
        self.prompted.set()

    async def list_messages_a(self, sid):
        return self.messages

    async def get_session_a(self, sid):
        return self.session_info

    async def abort_session_a(self, sid):
        self.aborted.append(sid)

    async def events(self):
        await self.prompted.wait()
        for ev in self.script:
            yield json.dumps(ev)
        await asyncio.sleep(3600)  # 진짜처럼 계속 열려 있음


@pytest.fixture
def fake(monkeypatch):
    def _install(script):
        oc = FakeOpencode(script)
        monkeypatch.setattr(stream_mod, "opencode_service", oc)
        monkeypatch.setattr(events_mod, "event_bus", events_mod.EventBus(source=oc.events))
        monkeypatch.setattr(stream_mod, "event_bus", events_mod.event_bus)
        return oc

    return _install


async def _login(client: AsyncClient, email, pw):
    res = await client.post("/api/v1/auth/login", {"email": email, "password": pw}, content_type="application/json")
    return res.json()["data"]["token"]


async def _collect(resp):
    """SSE 프레임 → [(event, data)]"""
    raw = b""
    async for chunk in resp.streaming_content:
        raw += chunk
    out = []
    for frame in raw.decode().split("\n\n"):
        if not frame.startswith("event:"):
            continue
        lines = frame.split("\n")
        ev = lines[0][7:]
        data = json.loads(lines[1][6:])
        out.append((ev, data))
    return out


ASSISTANT_DONE = [
    _msg_updated("mu", "user"),
    _part_updated("mu", "pu", "MARA 뭐야"),
    _msg_updated("ma", "assistant"),
    _part_updated("ma", "pa", "자재", delta="자재"),
    _part_updated("ma", "pa", "자재 마스터", delta=" 마스터"),
    {"type": "session.updated", "properties": {"info": {"id": "s1", "title": "MARA 설명"}}},
    {"type": "session.idle", "properties": {"sessionID": "s1"}},
]

FINAL_MESSAGES = [
    {"info": {"id": "mu", "role": "user"}, "parts": [{"type": "text", "text": "MARA 뭐야"}]},
    {
        "info": {
            "id": "ma",
            "role": "assistant",
            "tokens": {"input": 120, "output": 30, "reasoning": 0, "cache": {"read": 0, "write": 0}},
            "cost": 0.0012,
            "time": {"created": 1000, "completed": 3500},
        },
        "parts": [{"type": "text", "text": "자재 마스터"}],
    },
]


@pytest.mark.django_db(transaction=True)
async def test_stream_happy_path(fake, django_user_model):
    user = await django_user_model.objects.acreate(username="u@x.com", email="u@x.com", password="x")
    user.set_password("pw1234")
    await user.asave()
    await ChatSession.objects.acreate(id="s1", user=user)
    oc = fake(ASSISTANT_DONE)
    oc.messages = FINAL_MESSAGES

    client = AsyncClient()
    token = await _login(client, "u@x.com", "pw1234")
    resp = await client.post(
        "/api/v1/chat/stream",
        {"sessionId": "s1", "content": "MARA 뭐야", "explain": True},
        content_type="application/json",
        headers={"Authorization": f"Bearer {token}"},
    )
    assert resp.status_code == 200
    assert resp["Content-Type"].startswith("text/event-stream")
    frames = await _collect(resp)

    assert [f[0] for f in frames] == ["token", "token", "title", "usage", "done"]
    assert "".join(d["delta"] for e, d in frames if e == "token") == "자재 마스터"
    assert frames[2][1] == {"title": "MARA 설명"}
    assert frames[3][1] == {"used": 150, "limit": 128000, "ratio": round(150 / 128000, 4), "elapsed_ms": 2500}

    # OpenCode 에 보낸 payload — agent + explain 표시
    sid, payload = oc.payloads[0]
    assert sid == "s1" and payload["agent"] == "codeassist"
    assert payload["parts"] == [{"type": "text", "text": "[설명 모드] MARA 뭐야"}]

    # 미러 DB
    session = await ChatSession.objects.aget(id="s1")
    assert session.is_generating is False and session.title_llm == "MARA 설명"
    msgs = [m async for m in ChatMessage.objects.filter(session_id="s1").order_by("id")]
    assert [(m.role, m.content) for m in msgs] == [("user", "MARA 뭐야"), ("assistant", "자재 마스터")]
    assert (msgs[1].input_tokens, msgs[1].output_tokens, msgs[1].cost_usd, msgs[1].elapsed_ms) == (120, 30, 0.0012, 2500)


@pytest.mark.django_db(transaction=True)
async def test_stream_images_become_file_parts(fake, django_user_model):
    user = await django_user_model.objects.acreate(username="u@x.com", email="u@x.com", password="x")
    user.set_password("pw1234")
    await user.asave()
    await ChatSession.objects.acreate(id="s1", user=user)
    oc = fake([_msg_updated("ma", "assistant"), {"type": "session.idle", "properties": {"sessionID": "s1"}}])

    client = AsyncClient()
    token = await _login(client, "u@x.com", "pw1234")
    resp = await client.post(
        "/api/v1/chat/stream",
        {"sessionId": "s1", "content": "", "images": [{"mediaType": "image/png", "data": "data:image/png;base64,AA=="}]},
        content_type="application/json",
        headers={"Authorization": f"Bearer {token}"},
    )
    frames = await _collect(resp)
    assert frames[-1] == ("done", {})
    _, payload = oc.payloads[0]
    assert payload["parts"] == [
        {"type": "file", "mime": "image/png", "filename": "image-1.png", "url": "data:image/png;base64,AA=="}
    ]


@pytest.mark.django_db(transaction=True)
async def test_stream_session_error_emits_error_and_clears_generating(fake, django_user_model):
    user = await django_user_model.objects.acreate(username="u@x.com", email="u@x.com", password="x")
    user.set_password("pw1234")
    await user.asave()
    await ChatSession.objects.acreate(id="s1", user=user)
    fake(
        [
            _msg_updated("ma", "assistant"),
            _part_updated("ma", "pa", "절반", delta="절반"),
            {"type": "session.error", "properties": {"sessionID": "s1", "error": {"name": "ApiError", "data": {"message": "rate limited"}}}},
        ]
    )
    client = AsyncClient()
    token = await _login(client, "u@x.com", "pw1234")
    resp = await client.post(
        "/api/v1/chat/stream",
        {"sessionId": "s1", "content": "q"},
        content_type="application/json",
        headers={"Authorization": f"Bearer {token}"},
    )
    frames = await _collect(resp)
    assert frames == [("token", {"delta": "절반"}), ("error", {"message": "rate limited", "code": "LLM_ERROR"})]
    session = await ChatSession.objects.aget(id="s1")
    assert session.is_generating is False
    last = await ChatMessage.objects.filter(session_id="s1", role="assistant").alast()
    assert last.content == "절반"  # 부분 답변 보존


@pytest.mark.django_db(transaction=True)
async def test_stream_409_when_generating_and_401_without_token(fake, django_user_model):
    user = await django_user_model.objects.acreate(username="u@x.com", email="u@x.com", password="x")
    user.set_password("pw1234")
    await user.asave()
    await ChatSession.objects.acreate(id="s1", user=user, is_generating=True)
    fake([])
    client = AsyncClient()
    res = await client.post("/api/v1/chat/stream", {"sessionId": "s1", "content": "q"}, content_type="application/json")
    assert res.status_code == 401 and res.json()["code"] == "UNAUTHORIZED"
    token = await _login(client, "u@x.com", "pw1234")
    res = await client.post(
        "/api/v1/chat/stream",
        {"sessionId": "s1", "content": "q"},
        content_type="application/json",
        headers={"Authorization": f"Bearer {token}"},
    )
    assert res.status_code == 409 and res.json()["code"] == "CHAT_GENERATION_IN_PROGRESS"


@pytest.mark.django_db(transaction=True)
async def test_stream_validation(fake, django_user_model):
    user = await django_user_model.objects.acreate(username="u@x.com", email="u@x.com", password="x")
    user.set_password("pw1234")
    await user.asave()
    await ChatSession.objects.acreate(id="s1", user=user)
    fake([])
    client = AsyncClient()
    token = await _login(client, "u@x.com", "pw1234")
    hdr = {"Authorization": f"Bearer {token}"}
    res = await client.post("/api/v1/chat/stream", {"sessionId": "s1", "content": "  "}, content_type="application/json", headers=hdr)
    assert res.status_code == 400
    res = await client.post(
        "/api/v1/chat/stream",
        {"sessionId": "s1", "content": "x", "images": [{"mediaType": "image/gif", "data": "d"}]},
        content_type="application/json",
        headers=hdr,
    )
    assert res.status_code == 400
    res = await client.post("/api/v1/chat/stream", {"sessionId": "nope", "content": "x"}, content_type="application/json", headers=hdr)
    assert res.status_code == 404


@pytest.mark.django_db(transaction=True)
async def test_stream_no_events_at_all_fails_fast(fake, django_user_model, monkeypatch):
    """OpenCode 가 prompt 는 받았는데(204) 이벤트가 하나도 안 오면 → 첫 이벤트 타임아웃으로 error."""
    monkeypatch.setattr(stream_mod, "FIRST_EVENT_TIMEOUT_S", 0.2)
    monkeypatch.setattr(stream_mod, "KEEPALIVE_S", 0.05)
    user = await django_user_model.objects.acreate(username="u@x.com", email="u@x.com", password="x")
    user.set_password("pw1234")
    await user.asave()
    await ChatSession.objects.acreate(id="s1", user=user)
    oc = fake([])  # 이벤트 없음
    client = AsyncClient()
    token = await _login(client, "u@x.com", "pw1234")
    resp = await client.post(
        "/api/v1/chat/stream",
        {"sessionId": "s1", "content": "q"},
        content_type="application/json",
        headers={"Authorization": f"Bearer {token}"},
    )
    frames = await _collect(resp)
    assert frames[-1][0] == "error" and frames[-1][1]["code"] == "LLM_ERROR"
    assert oc.aborted == ["s1"]
    session = await ChatSession.objects.aget(id="s1")
    assert session.is_generating is False


@pytest.mark.django_db(transaction=True)
async def test_stream_title_arriving_after_idle_is_mirrored(fake, django_user_model, monkeypatch):
    monkeypatch.setattr(stream_mod, "TITLE_WAIT_S", 2)
    user = await django_user_model.objects.acreate(username="u@x.com", email="u@x.com", password="x")
    user.set_password("pw1234")
    await user.asave()
    await ChatSession.objects.acreate(id="s1", user=user)
    fake(
        [
            _msg_updated("ma", "assistant"),
            {"type": "session.updated", "properties": {"info": {"id": "s1", "title": "New session - 2026"}}},
            {"type": "session.idle", "properties": {"sessionID": "s1"}},
            {"type": "session.updated", "properties": {"info": {"id": "s1", "title": "늦게 온 제목"}}},
        ]
    )
    client = AsyncClient()
    token = await _login(client, "u@x.com", "pw1234")
    resp = await client.post(
        "/api/v1/chat/stream",
        {"sessionId": "s1", "content": "q"},
        content_type="application/json",
        headers={"Authorization": f"Bearer {token}"},
    )
    frames = await _collect(resp)
    assert frames[-1] == ("done", {})  # 제목 기다리느라 done 이 늦어지면 안 됨
    await asyncio.sleep(0.3)
    session = await ChatSession.objects.aget(id="s1")
    assert session.title_llm == "늦게 온 제목"