Aether/tests/api/handlers/base/test_stream_context.py

import os

import pytest

os.environ.setdefault("JWT_SECRET_KEY", "test-secret-key")

from src.api.handlers.base import stream_context
from src.api.handlers.base.stream_context import StreamContext


def test_collected_text_append_and_property() -> None:
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    assert ctx.collected_text == ""
    assert ctx.collected_text_length == 0

    ctx.append_text("hello")
    ctx.append_text(" ")
    ctx.append_text("world")
    assert ctx.collected_text == "hello world"
    assert ctx.collected_text_length == len("hello world")


def test_collected_text_is_capped_but_total_length_is_preserved() -> None:
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    cap = stream_context._MAX_COLLECTED_TEXT_CHARS

    ctx.append_text("a" * (cap - 4))
    ctx.append_text("b" * 10)

    assert len(ctx.collected_text) == cap
    assert ctx.collected_text == ("a" * (cap - 4)) + ("b" * 4)
    assert ctx.collected_text_length == cap + 6


def test_reset_for_retry_clears_state() -> None:
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    ctx.append_text("x")
    ctx.update_usage(input_tokens=10, output_tokens=5)
    ctx.parsed_chunks.append({"type": "chunk"})
    ctx.chunk_count = 3
    ctx.data_count = 2
    ctx.has_completion = True
    ctx.status_code = 418
    ctx.error_message = "boom"

    ctx.reset_for_retry()

    assert ctx.collected_text == ""
    assert ctx.input_tokens == 0
    assert ctx.output_tokens == 0
    assert ctx.parsed_chunks == []
    assert ctx.chunk_count == 0
    assert ctx.data_count == 0
    assert ctx.has_completion is False
    assert ctx.status_code == 200
    assert ctx.error_message is None


def test_release_recorded_chunks_clears_both_chunk_lists() -> None:
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    ctx.parsed_chunks.append({"type": "client"})
    ctx.provider_parsed_chunks.append({"type": "provider"})

    ctx.release_recorded_chunks()

    assert ctx.parsed_chunks == []
    assert ctx.provider_parsed_chunks == []


def test_managed_recorded_bodies_builds_then_releases_chunks() -> None:
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    ctx.parsed_chunks.append({"type": "client"})
    ctx.provider_parsed_chunks.append({"type": "provider"})
    ctx.data_count = 1

    with ctx.managed_recorded_bodies(123) as recorded_bodies:
        assert recorded_bodies.response_body is not None
        assert recorded_bodies.response_body["chunks"] == [{"type": "provider"}]
        assert recorded_bodies.client_response_body is not None
        assert recorded_bodies.client_response_body["chunks"] == [{"type": "client"}]

    assert ctx.parsed_chunks == []
    assert ctx.provider_parsed_chunks == []
    assert recorded_bodies.response_body is None
    assert recorded_bodies.client_response_body is None


def test_record_first_byte_time(monkeypatch: pytest.MonkeyPatch) -> None:
    """测试记录首字时间"""
    ctx = StreamContext(model="claude-3", api_format="claude_messages")
    start_time = 100.0
    monkeypatch.setattr(stream_context.time, "time", lambda: 100.0123)  # 12.3ms

    # 记录首字时间
    ctx.record_first_byte_time(start_time)

    # 验证首字时间已记录
    assert ctx.first_byte_time_ms == 12


def test_record_first_byte_time_idempotent(monkeypatch: pytest.MonkeyPatch) -> None:
    """测试首字时间只记录一次"""
    ctx = StreamContext(model="claude-3", api_format="claude_messages")
    start_time = 100.0

    # 第一次记录
    monkeypatch.setattr(stream_context.time, "time", lambda: 100.010)
    ctx.record_first_byte_time(start_time)
    first_value = ctx.first_byte_time_ms

    # 第二次记录（应该被忽略）
    monkeypatch.setattr(stream_context.time, "time", lambda: 100.020)
    ctx.record_first_byte_time(start_time)
    second_value = ctx.first_byte_time_ms

    # 验证值没有改变
    assert first_value == second_value


def test_reset_for_retry_clears_first_byte_time(monkeypatch: pytest.MonkeyPatch) -> None:
    """测试重试时清除首字时间"""
    ctx = StreamContext(model="claude-3", api_format="claude_messages")
    start_time = 100.0

    # 记录首字时间
    monkeypatch.setattr(stream_context.time, "time", lambda: 100.010)
    ctx.record_first_byte_time(start_time)
    assert ctx.first_byte_time_ms is not None

    # 重置
    ctx.reset_for_retry()

    # 验证首字时间已清除
    assert ctx.first_byte_time_ms is None


def test_get_log_summary_with_first_byte_time() -> None:
    """测试日志摘要包含首字时间"""
    ctx = StreamContext(model="claude-3", api_format="claude_messages")
    ctx.provider_name = "anthropic"
    ctx.input_tokens = 100
    ctx.output_tokens = 50
    ctx.first_byte_time_ms = 123

    summary = ctx.get_log_summary("request-id-123", 456)

    # 验证包含首字时间和总时间（大写格式）
    assert "TTFB: 123ms" in summary
    assert "Total: 456ms" in summary
    assert "in:100 out:50" in summary


def test_get_log_summary_without_first_byte_time() -> None:
    """测试日志摘要在没有首字时间时的格式"""
    ctx = StreamContext(model="claude-3", api_format="claude_messages")
    ctx.provider_name = "anthropic"
    ctx.input_tokens = 100
    ctx.output_tokens = 50
    # first_byte_time_ms 保持为 None

    summary = ctx.get_log_summary("request-id-123", 456)

    # 验证不包含首字时间标记,但有总时间（使用大写 TTFB 和 Total）
    assert "TTFB:" not in summary
    assert "Total: 456ms" in summary
    assert "in:100 out:50" in summary


def test_ensure_estimated_output_tokens_uses_collected_text() -> None:
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    ctx.append_text("partial output")

    changed = ctx.ensure_estimated_output_tokens()

    assert changed is True
    assert ctx.output_tokens == max(1, len("partial output") // 4)


def test_should_estimate_incomplete_tokens_for_interrupted_partial_stream() -> None:
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    ctx.status_code = 503
    ctx.chunk_count = 3

    assert ctx.should_estimate_incomplete_tokens() is True


def test_should_estimate_incomplete_tokens_when_output_already_estimated() -> None:
    """ensure_estimated_output_tokens 已补了 output，但 input 仍为 0 时仍需估算。"""
    ctx = StreamContext(model="test-model", api_format="openai:chat")
    ctx.status_code = 503
    ctx.chunk_count = 3
    ctx.append_text("partial")
    ctx.ensure_estimated_output_tokens()

    assert ctx.output_tokens > 0
    assert ctx.input_tokens == 0
    assert ctx.should_estimate_incomplete_tokens() is True