CaptchBreaker/tests/test_generators.py

"""
测试所有验证码生成器。

每种生成器 generate() 1 张 → 验证返回类型、图片尺寸、标签格式。
"""

import re

import pytest
from PIL import Image

from config import (
    GENERATE_CONFIG,
    NORMAL_CHARS,
    MATH_CHARS,
    THREED_CHARS,
    SOLVER_CONFIG,
    SOLVER_REGRESSION_RANGE,
)
from generators import (
    NormalCaptchaGenerator,
    MathCaptchaGenerator,
    ThreeDCaptchaGenerator,
    ThreeDRotateGenerator,
    ThreeDSliderGenerator,
    SlideDataGenerator,
    RotateSolverDataGenerator,
)


class TestNormalCaptchaGenerator:
    def setup_method(self):
        self.gen = NormalCaptchaGenerator(seed=0)
        self.cfg = GENERATE_CONFIG["normal"]

    def test_generate_returns_image_and_label(self):
        img, label = self.gen.generate()
        assert isinstance(img, Image.Image)
        assert isinstance(label, str)

    def test_image_size(self):
        img, _ = self.gen.generate()
        w, h = self.cfg["image_size"]
        assert img.size == (w, h)

    def test_label_chars_in_charset(self):
        img, label = self.gen.generate()
        assert len(label) >= 4
        for ch in label:
            assert ch in NORMAL_CHARS, f"char {ch!r} not in NORMAL_CHARS"

    def test_generate_with_text(self):
        img, label = self.gen.generate(text="AB12")
        assert label == "AB12"


class TestMathCaptchaGenerator:
    def setup_method(self):
        self.gen = MathCaptchaGenerator(seed=0)
        self.cfg = GENERATE_CONFIG["math"]

    def test_generate_returns_image_and_label(self):
        img, label = self.gen.generate()
        assert isinstance(img, Image.Image)
        assert isinstance(label, str)

    def test_image_size(self):
        img, _ = self.gen.generate()
        w, h = self.cfg["image_size"]
        assert img.size == (w, h)

    def test_label_is_expression(self):
        """Label should be like '3+8' (expression without =? and without result)."""
        img, label = self.gen.generate()
        assert re.match(r"^\d+[+\-×÷]\d+$", label), f"unexpected label format: {label!r}"

    def test_generate_with_division_text(self):
        img, label = self.gen.generate(text="20÷4")
        assert label == "20÷4"


class TestThreeDCaptchaGenerator:
    def setup_method(self):
        self.gen = ThreeDCaptchaGenerator(seed=0)
        self.cfg = GENERATE_CONFIG["3d_text"]

    def test_generate_returns_image_and_label(self):
        img, label = self.gen.generate()
        assert isinstance(img, Image.Image)
        assert isinstance(label, str)

    def test_image_size(self):
        img, _ = self.gen.generate()
        w, h = self.cfg["image_size"]
        assert img.size == (w, h)

    def test_label_chars_in_charset(self):
        img, label = self.gen.generate()
        assert len(label) >= 4
        for ch in label:
            assert ch in THREED_CHARS, f"char {ch!r} not in THREED_CHARS"


class TestThreeDRotateGenerator:
    def setup_method(self):
        self.gen = ThreeDRotateGenerator(seed=0)
        self.cfg = GENERATE_CONFIG["3d_rotate"]

    def test_generate_returns_image_and_label(self):
        img, label = self.gen.generate()
        assert isinstance(img, Image.Image)
        assert isinstance(label, str)

    def test_image_size(self):
        img, _ = self.gen.generate()
        w, h = self.cfg["image_size"]
        assert img.size == (w, h)

    def test_label_is_angle(self):
        img, label = self.gen.generate()
        angle = int(label)
        assert 0 <= angle <= 359


class TestThreeDSliderGenerator:
    def setup_method(self):
        self.gen = ThreeDSliderGenerator(seed=0)
        self.cfg = GENERATE_CONFIG["3d_slider"]

    def test_generate_returns_image_and_label(self):
        img, label = self.gen.generate()
        assert isinstance(img, Image.Image)
        assert isinstance(label, str)

    def test_image_size(self):
        img, _ = self.gen.generate()
        w, h = self.cfg["image_size"]
        assert img.size == (w, h)

    def test_label_is_offset(self):
        img, label = self.gen.generate()
        offset = int(label)
        lo, hi = self.cfg["gap_x_range"]
        assert lo <= offset <= hi


class TestSlideDataGenerator:
    def setup_method(self):
        self.gen = SlideDataGenerator(seed=0)

    def test_generate_returns_image_and_label(self):
        img, label = self.gen.generate()
        assert isinstance(img, Image.Image)
        assert isinstance(label, str)

    def test_image_size(self):
        img, _ = self.gen.generate()
        h, w = SOLVER_CONFIG["slide"]["cnn_input_size"]
        assert img.size == (w, h)

    def test_label_is_numeric(self):
        img, label = self.gen.generate()
        val = int(label)
        gs = self.gen.gap_size
        margin = gs + 10
        assert margin + gs // 2 <= val <= self.gen.width - margin + gs // 2

    def test_labels_normalize_inside_solver_range(self, tmp_path):
        for idx in range(3):
            img, label = self.gen.generate()
            img.save(tmp_path / f"{label}_{idx:06d}.png")

        from training.dataset import RegressionDataset

        ds = RegressionDataset(
            dirs=[tmp_path],
            label_range=SOLVER_REGRESSION_RANGE["slide"],
            transform=None,
        )
        assert len(ds.samples) == 3
        for _, norm in ds.samples:
            assert 0.0 < norm < 1.0


class TestRotateSolverDataGenerator:
    def setup_method(self):
        self.gen = RotateSolverDataGenerator(seed=0)

    def test_generate_returns_image_and_label(self):
        img, label = self.gen.generate()
        assert isinstance(img, Image.Image)
        assert isinstance(label, str)

    def test_image_size(self):
        img, _ = self.gen.generate()
        h, w = SOLVER_CONFIG["rotate"]["input_size"]
        assert img.size == (w, h)

    def test_label_is_angle(self):
        img, label = self.gen.generate()
        angle = int(label)
        assert 0 <= angle <= 359