llamacppctl/tests/test_http_ops.py

import sys
from pathlib import Path
from unittest.mock import patch, MagicMock

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import requests  # noqa: E402

from llamacppctl.http_ops import (  # noqa: E402
    HttpError,
    base_url,
    chat_completion,
    chat_completion_text,
    check_health,
    check_models,
    wait_until_ready,
)
from llamacppctl.schema import PromptConfig  # noqa: E402
from tests.test_docker_ops import make_cfg  # noqa: E402


def test_base_url():
    cfg = make_cfg(host_port=8001)
    assert base_url(cfg) == "http://127.0.0.1:8001/v1"


def test_check_health_ok():
    cfg = make_cfg()
    with patch("requests.get") as mock_get:
        mock_get.return_value = MagicMock(ok=True)
        assert check_health(cfg) is True


def test_check_health_exception_returns_false():
    cfg = make_cfg()
    with patch("requests.get", side_effect=requests.RequestException("boom")):
        assert check_health(cfg) is False


def test_check_models_ok():
    cfg = make_cfg()
    with patch("requests.get") as mock_get:
        mock_get.return_value = MagicMock(ok=True)
        assert check_models(cfg) is True


def test_chat_completion_posts_expected_payload():
    cfg = make_cfg()
    prompt_cfg = PromptConfig(system_prompt="sys", user_prompt="hi", max_tokens=32, temperature=0.1)
    with patch("requests.post") as mock_post:
        mock_post.return_value = MagicMock(status_code=200)
        chat_completion(cfg, prompt_cfg)
        _, kwargs = mock_post.call_args
        payload = kwargs["json"]
        assert payload["messages"][0] == {"role": "system", "content": "sys"}
        assert payload["messages"][1] == {"role": "user", "content": "hi"}
        assert payload["max_tokens"] == 32


def test_chat_completion_text_success():
    cfg = make_cfg()
    prompt_cfg = PromptConfig(user_prompt="hi")
    with patch("requests.post") as mock_post:
        mock_resp = MagicMock(status_code=200)
        mock_resp.json.return_value = {"choices": [{"message": {"content": "ok"}}]}
        mock_post.return_value = mock_resp
        reply = chat_completion_text(cfg, prompt_cfg)
        assert reply.content == "ok"
        assert reply.finish_reason == ""


def test_chat_completion_text_honors_chat_endpoint():
    cfg = make_cfg(host_port=8001, chat_endpoint="/custom/chat")
    prompt_cfg = PromptConfig(user_prompt="hi")
    with patch("requests.post") as mock_post:
        mock_resp = MagicMock(status_code=200)
        mock_resp.json.return_value = {"choices": [{"message": {"content": "ok"}}]}
        mock_post.return_value = mock_resp
        chat_completion_text(cfg, prompt_cfg)
        args, _ = mock_post.call_args
        assert args[0] == "http://127.0.0.1:8001/custom/chat"


def test_chat_completion_text_http_error():
    cfg = make_cfg()
    prompt_cfg = PromptConfig(user_prompt="hi")
    with patch("requests.post") as mock_post:
        mock_post.return_value = MagicMock(status_code=500, text="server error")
        try:
            chat_completion_text(cfg, prompt_cfg)
            assert False, "expected HttpError"
        except HttpError:
            pass


def test_chat_completion_text_transport_error():
    cfg = make_cfg()
    prompt_cfg = PromptConfig(user_prompt="hi")
    with patch("requests.post", side_effect=requests.RequestException("down")):
        try:
            chat_completion_text(cfg, prompt_cfg)
            assert False, "expected HttpError"
        except HttpError:
            pass


def test_chat_completion_text_malformed_response():
    cfg = make_cfg()
    prompt_cfg = PromptConfig(user_prompt="hi")
    with patch("requests.post") as mock_post:
        mock_resp = MagicMock(status_code=200)
        mock_resp.json.return_value = {"unexpected": True}
        mock_post.return_value = mock_resp
        try:
            chat_completion_text(cfg, prompt_cfg)
            assert False, "expected HttpError"
        except HttpError:
            pass


def test_chat_completion_text_reports_finish_reason():
    cfg = make_cfg()
    prompt_cfg = PromptConfig(user_prompt="hi")
    with patch("requests.post") as mock_post:
        mock_resp = MagicMock(status_code=200)
        mock_resp.json.return_value = {
            "choices": [{"finish_reason": "length", "message": {"content": ""}}]
        }
        mock_post.return_value = mock_resp
        reply = chat_completion_text(cfg, prompt_cfg)
        assert reply.content == ""
        assert reply.finish_reason == "length"


def test_api_key_sent_as_bearer():
    cfg = make_cfg(api_key="secret123")
    prompt_cfg = PromptConfig(user_prompt="hi")
    with patch("requests.post") as mock_post:
        mock_resp = MagicMock(status_code=200)
        mock_resp.json.return_value = {"choices": [{"message": {"content": "x"}}]}
        mock_post.return_value = mock_resp
        chat_completion_text(cfg, prompt_cfg)
        _, kwargs = mock_post.call_args
        assert kwargs["headers"]["Authorization"] == "Bearer secret123"


def test_stream_chat_yields_content_and_finish():
    cfg = make_cfg()
    prompt_cfg = PromptConfig(user_prompt="hi", stream=True)
    lines = [
        'data: {"choices":[{"delta":{"content":"Hal"}}]}',
        'data: {"choices":[{"delta":{"content":"lo"}}]}',
        'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}',
        "data: [DONE]",
    ]
    mock_resp = MagicMock(status_code=200)
    mock_resp.iter_lines.return_value = iter(lines)
    with patch("requests.post", return_value=mock_resp):
        from llamacppctl.http_ops import stream_chat

        events = list(stream_chat(cfg, prompt_cfg))
    contents = [v for k, v in events if k == "content"]
    finishes = [v for k, v in events if k == "finish"]
    assert "".join(contents) == "Hallo"
    assert finishes == ["stop"]
    # Stream must be forced to UTF-8 (llama.cpp sends no charset -> Umlaut-Bug).
    assert mock_resp.encoding == "utf-8"


def test_wait_until_ready_succeeds_immediately():
    cfg = make_cfg(timeout=5, poll_interval=0.01)
    probe = PromptConfig(user_prompt="ping", max_tokens=1)
    with patch("llamacppctl.http_ops.chat_completion") as mock_chat:
        mock_chat.return_value = MagicMock(status_code=200)
        assert wait_until_ready(cfg, probe) is True


def test_wait_until_ready_times_out():
    cfg = make_cfg(timeout=0, poll_interval=0.01)
    probe = PromptConfig(user_prompt="ping", max_tokens=1)
    with patch("llamacppctl.http_ops.chat_completion") as mock_chat:
        mock_chat.return_value = MagicMock(status_code=500)
        assert wait_until_ready(cfg, probe) is False