Spaces:
Running
Running
File size: 2,410 Bytes
e6404d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | 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 |