"""
单元测试 — 模型模块 (Person B 实现后需通过)
"""

import pytest
import torch


class TestMultiHeadAttention:
    """测试多头注意力。"""

    def test_output_shape(self):
        """测试输出维度正确。"""
        # TODO: 创建 MHA，输入 [2, 10, 512]，验证输出 [2, 10, 512]
        pass

    def test_padding_mask(self):
        """测试 padding mask 效果。"""
        pass


class TestPositionalEncoding:
    """测试位置编码。"""

    def test_sinusoidal_shape(self):
        """测试正弦位置编码输出维度。"""
        pass

    def test_rope_shape(self):
        """测试 RoPE 输出维度。"""
        pass

    def test_rope_relative_distance(self):
        """测试 RoPE 捕获相对距离。"""
        pass


class TestTransformerModel:
    """测试完整 Transformer 模型。"""

    def test_forward_shape(self):
        """测试前向传播输出维度。"""
        # TODO: 构建小模型，输入 mock 数据，验证 logits shape [B, T, V]
        pass

    def test_encode_shape(self):
        """测试 encode 输出维度。"""
        pass

    def test_decode_step_shape(self):
        """测试 decode_step 输出维度。"""
        pass

    def test_parameter_count(self):
        """测试参数数量计算。"""
        pass

    def test_causal_mask(self):
        """测试因果掩码生成。"""
        pass