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