vibe-bot/vibe_bot/tests/test_prompts.py

"""Tests for the prompts module (prompt constants, layout parsing, user info)."""

from __future__ import annotations

from datetime import UTC, datetime
from unittest.mock import MagicMock

import pytest

from vibe_bot.config import (
    IMAGE_GEN_SIZE_LANDSCAPE,
    IMAGE_GEN_SIZE_PORTRAIT,
    IMAGE_GEN_SIZE_SQUARE,
)
from vibe_bot.prompts import (
    IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE,
    LAYOUT_SIZES,
    RESPONSE_LENGTH_HINT,
    build_system_prompt,
    get_user_info,
    parse_image_layout,
)


@pytest.fixture
def mock_ctx() -> MagicMock:
    """A minimal Discord user for get_user_info."""
    author = MagicMock()
    author.name = "testuser"
    author.id = "12345"
    author.global_name = "Test User"
    author.nick = "tester"
    author.top_role.name = "@everyone"
    author.activities = []
    author.joined_at = None
    author.created_at = None
    return author


@pytest.fixture
def mock_author_with_member_data() -> MagicMock:
    """A Discord user with full member data (role + activity + timestamps)."""
    author = MagicMock()
    author.name = "testuser"
    author.id = "12345"
    author.global_name = "Test User"
    author.nick = "tester"
    author.top_role.name = "Admin"
    activity = MagicMock()
    activity.name = "Chess"
    author.activities = [activity]
    author.joined_at = datetime(2024, 1, 15, tzinfo=UTC)
    author.created_at = datetime(2023, 6, 1, tzinfo=UTC)
    return author


def test_build_system_prompt_contains_personality_hint_and_user_info() -> None:
    """build_system_prompt assembles personality + length hint + user info block."""
    result = build_system_prompt("you are a butler", "Username: alice")
    assert result.startswith("you are a butler")
    assert RESPONSE_LENGTH_HINT in result
    assert "User Information:\nUsername: alice" in result


def test_layout_sizes_map_to_config() -> None:
    """LAYOUT_SIZES maps each layout to its configured canvas size."""
    assert LAYOUT_SIZES["portrait"] == IMAGE_GEN_SIZE_PORTRAIT
    assert LAYOUT_SIZES["landscape"] == IMAGE_GEN_SIZE_LANDSCAPE
    assert LAYOUT_SIZES["square"] == IMAGE_GEN_SIZE_SQUARE


@pytest.mark.parametrize(
    ("response", "expected"),
    [
        ("portrait", "portrait"),
        ("landscape", "landscape"),
        ("square", "square"),
        ("  Portrait  ", "portrait"),
        ("LANDSCAPE", "landscape"),
        ("square.", "square"),
        ("I would use portrait.", "portrait"),
        ("This scene is best as landscape.", "landscape"),
        ("A square composition works here.", "square"),
    ],
)
def test_parse_image_layout_valid(response: str, expected: str) -> None:
    """Valid LLM layout responses parse to the right layout."""
    assert parse_image_layout(response) == expected


@pytest.mark.parametrize(
    "response",
    [
        "",
        "   ",
        "banana",
        "1024x1024",
        "tall and wide",
        "squareness",  # word boundary should prevent a match on "square"
    ],
)
def test_parse_image_layout_defaults_to_square(response: str) -> None:
    """Empty or malformed responses fall back to square."""
    assert parse_image_layout(response) == "square"


def test_image_prompt_system_prompt_covers_key_details() -> None:
    """The prompt-rewrite system prompt forces explicit detail on all aspects."""
    prompt = IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE.format(layout="square")
    lowered = prompt.lower()
    assert "square" in prompt
    assert "exact text" in lowered
    assert "composition" in lowered
    assert "style" in lowered
    assert "horrifying creature" in lowered
    assert "only the image generation prompt" in lowered
    assert "exactly as written" in lowered
    assert "fountain pen wearing pants" in lowered
    assert "centaur" in lowered
    assert "not a person riding a horse" in lowered


def test_get_user_info_minimal(mock_ctx: MagicMock) -> None:
    """get_user_info with minimal member data includes the core identity lines."""
    result = get_user_info(mock_ctx)
    assert "Username: testuser" in result
    assert "User ID: 12345" in result
    assert "Global Name: Test User" in result
    assert "Nickname: tester" in result


def test_get_user_info_with_member_data(
    mock_author_with_member_data: MagicMock,
) -> None:
    """get_user_info with full member data includes roles, activity, timestamps."""
    result = get_user_info(mock_author_with_member_data)
    assert "Global Name: Test User" in result
    assert "Nickname: tester" in result
    assert "Username: testuser" in result
    assert "User ID: 12345" in result
    assert "Top Role: Admin" in result
    assert "Activities: Chess" in result
    assert "Joined: 2024-01-15" in result
    assert "Account Created: 2023-06-01" in result


def test_get_user_info_no_global_name(mock_ctx: MagicMock) -> None:
    """Optional fields are omitted when they are empty."""
    mock_ctx.global_name = None
    mock_ctx.nick = None
    mock_ctx.top_role.name = "@everyone"
    mock_ctx.activities = []

    result = get_user_info(mock_ctx)

    assert "Global Name:" not in result
    assert "Nickname:" not in result
    assert "Top Role:" not in result
    assert "Activities:" not in result
    assert "Username: testuser" in result
    assert "User ID: 12345" in result


def test_get_user_info_with_top_role_not_everyone(
    mock_author_with_member_data: MagicMock,
) -> None:
    """Top role is included when it is not @everyone."""
    result = get_user_info(mock_author_with_member_data)
    assert "Top Role: Admin" in result


def test_get_user_info_no_activities(mock_ctx: MagicMock) -> None:
    """The activities line is omitted when there are none."""
    mock_ctx.activities = []
    result = get_user_info(mock_ctx)
    assert "Activities:" not in result