72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from moderation import detect_person, is_prompt_allowed
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("prompt", "mode", "has_person", "reason"),
|
|
[
|
|
("child nude", "t2i", False, "csam"),
|
|
("l0li", "t2i", False, "csam"),
|
|
("sex with dog", "t2i", False, "bestiality"),
|
|
("Trump nude", "t2i", False, "real_person"),
|
|
("nude woman", "t2i", False, "nonconsensual_sexualization"),
|
|
("голая женщина", "t2i", False, "nonconsensual_sexualization"),
|
|
("blood sex", "t2i", False, "violence_sexual"),
|
|
("how to make a bomb", "t2i", False, "extremism"),
|
|
("nazi symbol", "t2i", False, "extremism"),
|
|
("ch1ld", "t2i", False, "obfuscated_nsfw"),
|
|
("nude", "i2i", True, "i2i_person_nsfw"),
|
|
],
|
|
)
|
|
def test_blocked_prompts(
|
|
prompt: str, mode: str, has_person: bool, reason: str
|
|
) -> None:
|
|
allowed, actual_reason = is_prompt_allowed(
|
|
prompt, mode=mode, has_person=has_person
|
|
)
|
|
assert not allowed
|
|
assert actual_reason == reason
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"prompt",
|
|
[
|
|
"red apple on a wooden table, soft light",
|
|
"nude beach landscape",
|
|
"cat sitting on a chair",
|
|
"голубое небо, женщина на пляже",
|
|
],
|
|
)
|
|
def test_allowed_prompts(prompt: str) -> None:
|
|
allowed, reason = is_prompt_allowed(prompt)
|
|
assert allowed
|
|
assert reason is None
|
|
|
|
|
|
def test_admin_bypasses_filters() -> None:
|
|
allowed, reason = is_prompt_allowed("child nude", is_admin=True)
|
|
assert allowed
|
|
assert reason is None
|
|
|
|
|
|
def test_extra_public_figures() -> None:
|
|
allowed, reason = is_prompt_allowed(
|
|
"customperson nude",
|
|
extra_public_figures={"customperson"},
|
|
)
|
|
assert not allowed
|
|
assert reason == "real_person"
|
|
|
|
|
|
def test_detect_person_blank_image(tmp_path) -> None:
|
|
pytest.importorskip("cv2")
|
|
path = tmp_path / "blank.png"
|
|
with Image.new("RGB", (320, 240), (255, 255, 255)) as image:
|
|
image.save(path)
|
|
assert detect_person(path) is False
|
|
assert detect_person(tmp_path / "missing.png") is False
|