Function generate_test_class

utils/generate_model_tests.py:313–429

(model_name: str, config_class: str, tester: str)

def generate_test_class(model_name: str, config_class: str, tester: str) -> str:
    tester_short = tester.replace("TesterMixin", "")
    class_name = f"Test{model_name}{tester_short}"

    lines = [f"class {class_name}({config_class}, {tester}):"]

    if tester == "TorchCompileTesterMixin":
        lines.extend(
            [
                "    @property",
                "    def different_shapes_for_compilation(self):",
                "        return [(4, 4), (4, 8), (8, 8)]",
                "",
                "    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:",
                "        # TODO: Implement dynamic input generation",
                "        return {}",
            ]
        )
    elif tester == "IPAdapterTesterMixin":
        lines.extend(
            [
                "    @property",
                "    def ip_adapter_processor_cls(self):",
                "        return None  # TODO: Set processor class",
                "",
                "    def modify_inputs_for_ip_adapter(self, model, inputs_dict):",
                "        # TODO: Add IP adapter image embeds to inputs",
                "        return inputs_dict",
                "",
                "    def create_ip_adapter_state_dict(self, model):",
                "        # TODO: Create IP adapter state dict",
                "        return {}",
            ]
        )
    elif tester == "SingleFileTesterMixin":
        lines.extend(
            [
                "    @property",
                "    def ckpt_path(self):",
                '        return ""  # TODO: Set checkpoint path',
                "",
                "    @property",
                "    def alternate_ckpt_paths(self):",
                "        return []",
                "",
                "    @property",
                "    def pretrained_model_name_or_path(self):",
                '        return ""  # TODO: Set Hub repository ID',
            ]
        )
    elif tester == "GGUFTesterMixin":
        lines.extend(
            [
                "    @property",
                "    def gguf_filename(self):",
                '        return ""  # TODO: Set GGUF filename',
                "",
                "    def get_dummy_inputs(self) -> dict[str, torch.Tensor]:",