tests/test_handlers.py

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import pytest
from src.handlers import (
    on_generate, on_prompt_analyze, on_user_preset_save,
    on_user_preset_apply, on_user_preset_delete,
    on_web_search, on_rating_change, on_artist_filter_change,
    on_tagger_apply_tags, MAX_OUTPUTS,
    ALL_CATEGORIES,
)
from src.variation_engine import _estimate_tokens


def test_all_categories_list_exists():
    assert len(ALL_CATEGORIES) > 0
    assert "quality" in ALL_CATEGORIES
    assert "nsfw" in ALL_CATEGORIES


def test_on_generate_returns_correct_length():
    n_checks = len(ALL_CATEGORIES)
    result = on_generate(
        "1girl, blue hair", "anima", "pg", 2, "medium", "off",
        "normal", "en", "", [], False, None, "0", "", "", "off",
        "", False, False, False, False, 0, "", "prompt",
        *([True] * n_checks),
    )
    assert len(result) == MAX_OUTPUTS * 2


def test_on_generate_empty_returns_empty_tuple():
    n_checks = len(ALL_CATEGORIES)
    result = on_generate(
        "", "anima", "pg", 3, "medium", "off",
        "normal", "en", "", [], False, None, "0", "", "", "off",
        "", False, False, False, False, 0, "", "prompt",
        *([True] * n_checks),
    )
    assert len(result) == MAX_OUTPUTS * 2
    assert all(x == "" for x in result)


def test_on_rating_change_pg_disables_nsfw():
    n_checks = len(ALL_CATEGORIES)
    checks = [True] * n_checks
    result = on_rating_change("pg", *checks)
    nsfw_idx = ALL_CATEGORIES.index("nsfw")
    assert result[nsfw_idx] is False


def test_on_rating_change_r_allows_nsfw():
    n_checks = len(ALL_CATEGORIES)
    checks = [True] * n_checks
    result = on_rating_change("r", *checks)
    nsfw_idx = ALL_CATEGORIES.index("nsfw")
    assert result[nsfw_idx] is True


def test_on_tagger_apply_tags():
    assert on_tagger_apply_tags("", "tag1, tag2") == "tag1, tag2"
    assert on_tagger_apply_tags("prompt", "tag1") == "prompt, tag1"
    assert on_tagger_apply_tags("prompt", "") == "prompt"
    assert on_tagger_apply_tags("", "") == ""


def test_estimate_tokens():
    est = _estimate_tokens("masterpiece, best quality, 1girl, blue hair, smile")
    assert isinstance(est, int)
    assert est > 0


def test_handlers_importable():
    assert on_generate is not None
    assert on_prompt_analyze is not None
    assert on_tagger_apply_tags is not None